Merge "Fix jank resulting from TaskView resizing" into main
diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index c3c42fa..aafa1f6 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -65,6 +65,13 @@
 }
 
 flag {
+    name: "enable_taskbar_connected_displays"
+    namespace: "launcher"
+    description: "Enables connected displays in taskbar."
+    bug: "362720616"
+}
+
+flag {
     name: "enable_taskbar_customization"
     namespace: "launcher"
     description: "Enables taskbar customization framework."
@@ -441,6 +448,12 @@
     bug: "292269949"
 }
 
+flag {
+    name: "enable_state_manager_proto_log"
+    namespace: "launcher"
+    description: "Enables tracking state manager logs in ProtoLog"
+    bug: "292269949"
+}
 
 flag {
     name: "coordinate_workspace_scale"
diff --git a/quickstep/res/layout/customizable_taskbar.xml b/quickstep/res/layout/customizable_taskbar.xml
index e1a80ae..d988cbc 100644
--- a/quickstep/res/layout/customizable_taskbar.xml
+++ b/quickstep/res/layout/customizable_taskbar.xml
@@ -51,20 +51,26 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent"/>
 
-    <com.android.launcher3.taskbar.bubbles.BubbleBarView
-        android:id="@+id/taskbar_bubbles"
-        android:layout_width="wrap_content"
-        android:layout_height="@dimen/bubblebar_size_with_pointer"
-        android:layout_gravity="bottom|end"
-        android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
-        android:paddingTop="@dimen/bubblebar_pointer_visible_size"
-        android:paddingEnd="@dimen/taskbar_icon_spacing"
-        android:paddingStart="@dimen/taskbar_icon_spacing"
-        android:visibility="gone"
-        android:gravity="center"
-        android:clipChildren="false"
-        android:elevation="@dimen/bubblebar_elevation"
-        />
+    <FrameLayout
+        android:id="@+id/taskbar_bubbles_container"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:clipChildren="false">
+
+        <com.android.launcher3.taskbar.bubbles.BubbleBarView
+            android:id="@+id/taskbar_bubbles"
+            android:layout_width="wrap_content"
+            android:layout_height="@dimen/bubblebar_size_with_pointer"
+            android:layout_gravity="bottom|end"
+            android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
+            android:paddingTop="@dimen/bubblebar_pointer_visible_size"
+            android:paddingEnd="@dimen/taskbar_icon_spacing"
+            android:paddingStart="@dimen/taskbar_icon_spacing"
+            android:visibility="gone"
+            android:gravity="center"
+            android:clipChildren="false"
+            android:elevation="@dimen/bubblebar_elevation" />
+    </FrameLayout>
 
     <com.android.launcher3.taskbar.navbutton.NearestTouchFrame
         android:id="@+id/navbuttons_view"
diff --git a/quickstep/res/layout/taskbar.xml b/quickstep/res/layout/taskbar.xml
index e8f3d9d..54f9ae8 100644
--- a/quickstep/res/layout/taskbar.xml
+++ b/quickstep/res/layout/taskbar.xml
@@ -35,17 +35,24 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent"/>
 
-    <com.android.launcher3.taskbar.bubbles.BubbleBarView
-        android:id="@+id/taskbar_bubbles"
-        android:layout_width="wrap_content"
-        android:layout_height="@dimen/bubblebar_size_with_pointer"
-        android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
-        android:paddingTop="@dimen/bubblebar_pointer_visible_size"
-        android:visibility="gone"
-        android:gravity="center"
-        android:layout_gravity="bottom"
-        android:clipChildren="false"
-        android:elevation="@dimen/bubblebar_elevation" />
+    <FrameLayout
+        android:id="@+id/taskbar_bubbles_container"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:clipChildren="false">
+
+        <com.android.launcher3.taskbar.bubbles.BubbleBarView
+            android:id="@+id/taskbar_bubbles"
+            android:layout_width="wrap_content"
+            android:layout_height="@dimen/bubblebar_size_with_pointer"
+            android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
+            android:paddingTop="@dimen/bubblebar_pointer_visible_size"
+            android:visibility="gone"
+            android:gravity="center"
+            android:layout_gravity="bottom"
+            android:clipChildren="false"
+            android:elevation="@dimen/bubblebar_elevation" />
+    </FrameLayout>
 
     <com.android.launcher3.taskbar.navbutton.NearestTouchFrame
         android:id="@+id/navbuttons_view"
diff --git a/quickstep/res/layout/transient_taskbar.xml b/quickstep/res/layout/transient_taskbar.xml
index f3c3383..3ec8046 100644
--- a/quickstep/res/layout/transient_taskbar.xml
+++ b/quickstep/res/layout/transient_taskbar.xml
@@ -38,18 +38,24 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent"/>
 
-    <com.android.launcher3.taskbar.bubbles.BubbleBarView
-        android:id="@+id/taskbar_bubbles"
-        android:layout_width="wrap_content"
-        android:layout_height="@dimen/bubblebar_size_with_pointer"
-        android:layout_gravity="bottom|end"
-        android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
-        android:paddingTop="@dimen/bubblebar_pointer_visible_size"
-        android:visibility="gone"
-        android:gravity="center"
-        android:clipChildren="false"
-        android:elevation="@dimen/bubblebar_elevation"
-        />
+    <FrameLayout
+        android:id="@+id/taskbar_bubbles_container"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:clipChildren="false">
+
+        <com.android.launcher3.taskbar.bubbles.BubbleBarView
+            android:id="@+id/taskbar_bubbles"
+            android:layout_width="wrap_content"
+            android:layout_height="@dimen/bubblebar_size_with_pointer"
+            android:layout_gravity="bottom|end"
+            android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
+            android:paddingTop="@dimen/bubblebar_pointer_visible_size"
+            android:visibility="gone"
+            android:gravity="center"
+            android:clipChildren="false"
+            android:elevation="@dimen/bubblebar_elevation" />
+    </FrameLayout>
 
     <com.android.launcher3.taskbar.navbutton.NearestTouchFrame
         android:id="@+id/navbuttons_view"
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index 6af5a30..4014f06 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -237,7 +237,7 @@
             InstanceId instanceId = new InstanceIdSequence().newInstanceId();
             for (ItemInfo info : itemsIdMap) {
                 CollectionInfo parent = getContainer(info, itemsIdMap);
-                StatsLogCompatManager.writeSnapshot(info.buildProto(parent), instanceId);
+                StatsLogCompatManager.writeSnapshot(info.buildProto(parent, mContext), instanceId);
             }
             additionalSnapshotEvents(instanceId);
             prefs.put(LAST_SNAPSHOT_TIME_MILLIS, now);
@@ -274,7 +274,7 @@
 
                         for (ItemInfo info : itemsIdMap) {
                             CollectionInfo parent = getContainer(info, itemsIdMap);
-                            LauncherAtom.ItemInfo itemInfo = info.buildProto(parent);
+                            LauncherAtom.ItemInfo itemInfo = info.buildProto(parent, mContext);
                             Log.d(TAG, itemInfo.toString());
                             StatsEvent statsEvent = StatsLogCompatManager.buildStatsEvent(itemInfo,
                                     instanceId);
diff --git a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
index 3dcb2ac..2ac87ff 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
@@ -16,7 +16,7 @@
 package com.android.launcher3.statehandlers;
 
 import static android.view.View.VISIBLE;
-import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY;
+import static android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY;
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index 10ff9ac..09dbeb6 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -15,8 +15,10 @@
  */
 package com.android.launcher3.taskbar;
 
-import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY;
+import static android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY;
 
+import static com.android.launcher3.QuickstepTransitionManager.TASKBAR_TO_APP_DURATION;
+import static com.android.launcher3.QuickstepTransitionManager.getTaskbarToHomeDuration;
 import static com.android.launcher3.QuickstepTransitionManager.TRANSIENT_TASKBAR_TRANSITION_DURATION;
 import static com.android.launcher3.statemanager.BaseState.FLAG_NON_INTERACTIVE;
 import static com.android.launcher3.taskbar.TaskbarEduTooltipControllerKt.TOOLTIP_STEP_FEATURES;
@@ -32,7 +34,6 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Flags;
 import com.android.launcher3.LauncherState;
-import com.android.launcher3.QuickstepTransitionManager;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimatedFloat;
 import com.android.launcher3.logging.InstanceId;
@@ -205,11 +206,17 @@
                 isVisible,
                 fromInitOrDestroy,
                 /* startAnimation= */ true,
-                DisplayController.isTransientTaskbar(mLauncher)
-                        ? TRANSIENT_TASKBAR_TRANSITION_DURATION
-                        : (!isVisible
-                                ? QuickstepTransitionManager.TASKBAR_TO_APP_DURATION
-                                : QuickstepTransitionManager.getTaskbarToHomeDuration()));
+                getTaskbarAnimationDuration(isVisible));
+    }
+
+    private int getTaskbarAnimationDuration(boolean isVisible) {
+        if (isVisible && !mLauncher.getPredictiveBackToHomeInProgress()) {
+            return getTaskbarToHomeDuration();
+        } else {
+            return DisplayController.isTransientTaskbar(mLauncher)
+                    ? TRANSIENT_TASKBAR_TRANSITION_DURATION
+                    : TASKBAR_TO_APP_DURATION;
+        }
     }
 
     @Nullable
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index d1725bc..7d8e93c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -173,9 +173,9 @@
     // Used for IME+A11Y buttons
     private final ViewGroup mEndContextualContainer;
     private final ViewGroup mStartContextualContainer;
-    private final int mLightIconColorOnHome;
-    private final int mDarkIconColorOnHome;
-    /** Color to use for navigation bar buttons, if they are on on a Taskbar surface background. */
+    private final int mLightIconColorOnWorkspace;
+    private final int mDarkIconColorOnWorkspace;
+    /** Color to use for navbar buttons, if they are on on a Taskbar surface background. */
     private final int mOnBackgroundIconColor;
 
     private @Nullable Animator mNavBarLocationAnimator;
@@ -191,7 +191,10 @@
     // Used for System UI state updates that should translate the nav button for in-app display.
     private final AnimatedFloat mNavButtonInAppDisplayProgressForSysui = new AnimatedFloat(
             this::updateNavButtonInAppDisplayProgressForSysui);
-    /** Expected nav button dark intensity communicated via the framework. */
+    /**
+     * Expected nav button dark intensity piped down from {@code LightBarController} in framework
+     * via {@code TaskbarDelegate}.
+     */
     private final AnimatedFloat mTaskbarNavButtonDarkIntensity = new AnimatedFloat(
             this::onDarkIntensityChanged);
     /** {@code 1} if the Taskbar background color is fully opaque. */
@@ -246,8 +249,8 @@
         mEndContextualContainer = mNavButtonsView.findViewById(R.id.end_contextual_buttons);
         mStartContextualContainer = mNavButtonsView.findViewById(R.id.start_contextual_buttons);
 
-        mLightIconColorOnHome = context.getColor(R.color.taskbar_nav_icon_light_color_on_home);
-        mDarkIconColorOnHome = context.getColor(R.color.taskbar_nav_icon_dark_color_on_home);
+        mLightIconColorOnWorkspace = context.getColor(R.color.taskbar_nav_icon_light_color_on_home);
+        mDarkIconColorOnWorkspace = context.getColor(R.color.taskbar_nav_icon_dark_color_on_home);
         mOnBackgroundIconColor = Utilities.isDarkTheme(context)
                 ? context.getColor(R.color.taskbar_nav_icon_light_color)
                 : context.getColor(R.color.taskbar_nav_icon_dark_color);
@@ -762,40 +765,68 @@
         mNavButtonsView.setTranslationY(mLastSetNavButtonTranslationY);
     }
 
+    /**
+     * Sets Taskbar 3-button mode icon colors based on the
+     * {@link #mTaskbarNavButtonDarkIntensity} value piped in from Framework. For certain cases
+     * in large screen taskbar where there may be opaque surfaces, the selected SystemUI button
+     * colors are intentionally overridden.
+     * <p>
+     * This method is also called when any of the AnimatedFloat instances change.
+     */
     private void updateNavButtonColor() {
         final ArgbEvaluator argbEvaluator = ArgbEvaluator.getInstance();
-        final int sysUiNavButtonIconColorOnHome = (int) argbEvaluator.evaluate(
-                mTaskbarNavButtonDarkIntensity.value,
-                mLightIconColorOnHome,
-                mDarkIconColorOnHome);
-
-        final int iconColor;
-        if (ENABLE_TASKBAR_NAVBAR_UNIFICATION && mContext.isPhoneMode()) {
-            iconColor = sysUiNavButtonIconColorOnHome;
-        } else {
-            // Override the color from framework if nav buttons are over an opaque Taskbar surface.
-            iconColor = (int) argbEvaluator.evaluate(
-                    mOnBackgroundNavButtonColorOverrideMultiplier.value * Math.max(
-                            mOnTaskbarBackgroundNavButtonColorOverride.value,
-                            mSlideInViewVisibleNavButtonColorOverride.value),
-                    sysUiNavButtonIconColorOnHome,
-                    mOnBackgroundIconColor);
+        int taskbarNavButtonColor = getSysUiIconColorOnHome(argbEvaluator);
+        // Only phone mode foldable button colors should be identical to SysUI navbar colors.
+        if (!(ENABLE_TASKBAR_NAVBAR_UNIFICATION && mContext.isPhoneMode())) {
+            taskbarNavButtonColor = getTaskbarButtonColor(argbEvaluator, taskbarNavButtonColor);
         }
+        applyButtonColors(taskbarNavButtonColor);
+    }
 
+    /**
+     * Taskbar 3-button mode icon colors based on the
+     * {@link #mTaskbarNavButtonDarkIntensity} value piped in from Framework.
+     */
+    private int getSysUiIconColorOnHome(ArgbEvaluator argbEvaluator) {
+        return (int) argbEvaluator.evaluate(getTaskbarNavButtonDarkIntensity().value,
+                mLightIconColorOnWorkspace, mDarkIconColorOnWorkspace);
+    }
+
+    /**
+     * If Taskbar background is opaque or slide in overlay is visible, the selected SystemUI button
+     * colors are intentionally overridden. The override can be disabled when
+     * {@link #mOnBackgroundNavButtonColorOverrideMultiplier} is {@code 0}.
+     */
+    private int getTaskbarButtonColor(ArgbEvaluator argbEvaluator, int sysUiIconColorOnHome) {
+        final float sysUIColorOverride =
+                mOnBackgroundNavButtonColorOverrideMultiplier.value * Math.max(
+                        mOnTaskbarBackgroundNavButtonColorOverride.value,
+                        mSlideInViewVisibleNavButtonColorOverride.value);
+        return (int) argbEvaluator.evaluate(sysUIColorOverride, sysUiIconColorOnHome,
+                mOnBackgroundIconColor);
+    }
+
+    /**
+     * Iteratively sets button colors for each button in {@link #mAllButtons}.
+     */
+    private void applyButtonColors(int iconColor) {
         for (ImageView button : mAllButtons) {
             button.setImageTintList(ColorStateList.valueOf(iconColor));
             Drawable background = button.getBackground();
             if (background instanceof KeyButtonRipple) {
                 ((KeyButtonRipple) background).setDarkIntensity(
-                        mTaskbarNavButtonDarkIntensity.value);
+                        getTaskbarNavButtonDarkIntensity().value);
             }
         }
     }
 
+    /**
+     * Updates Taskbar 3-Button icon colors as {@link #mTaskbarNavButtonDarkIntensity} changes.
+     */
     private void onDarkIntensityChanged() {
         updateNavButtonColor();
         if (mContext.isPhoneMode()) {
-            mTaskbarTransitions.onDarkIntensityChanged(mTaskbarNavButtonDarkIntensity.value);
+            mTaskbarTransitions.onDarkIntensityChanged(getTaskbarNavButtonDarkIntensity().value);
         }
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index f9efea9..1b9614a 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -71,6 +71,7 @@
 import android.view.View;
 import android.view.WindowInsets;
 import android.view.WindowManager;
+import android.widget.FrameLayout;
 import android.widget.Toast;
 import android.window.RemoteTransition;
 
@@ -268,8 +269,10 @@
         NearestTouchFrame navButtonsView = mDragLayer.findViewById(R.id.navbuttons_view);
         StashedHandleView stashedHandleView = mDragLayer.findViewById(R.id.stashed_handle);
         BubbleBarView bubbleBarView = null;
+        FrameLayout bubbleBarContainer = null;
         if (isTransientTaskbar || Flags.enableBubbleBarInPersistentTaskBar()) {
             bubbleBarView = mDragLayer.findViewById(R.id.taskbar_bubbles);
+            bubbleBarContainer = mDragLayer.findViewById(R.id.taskbar_bubbles_container);
         }
         StashedHandleView bubbleHandleView = mDragLayer.findViewById(R.id.stashed_bubble_handle);
 
@@ -296,7 +299,7 @@
                     : new PersistentBubbleStashController(dimensionsProvider);
             bubbleControllersOptional = Optional.of(new BubbleControllers(
                     new BubbleBarController(this, bubbleBarView),
-                    new BubbleBarViewController(this, bubbleBarView),
+                    new BubbleBarViewController(this, bubbleBarView, bubbleBarContainer),
                     bubbleStashController,
                     bubbleHandleController,
                     new BubbleDragController(this),
@@ -987,8 +990,8 @@
     }
 
     public void onNavButtonsDarkIntensityChanged(float darkIntensity) {
-        mControllers.navbarButtonsViewController.getTaskbarNavButtonDarkIntensity()
-                .updateValue(darkIntensity);
+        mControllers.navbarButtonsViewController.getTaskbarNavButtonDarkIntensity().updateValue(
+                darkIntensity);
     }
 
     public void onNavigationBarLumaSamplingEnabled(int displayId, boolean enable) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
index 219a24a..4a85acc 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -195,11 +195,12 @@
         };
 
         if (taskbarDesktopModeController.getAreDesktopTasksVisible()) {
-            mCornerRoundness.updateValue(taskbarDesktopModeController.getTaskbarCornerRoundness(
-                    mSharedState.showCornerRadiusInDesktopMode));
+            mCornerRoundness.value = taskbarDesktopModeController.getTaskbarCornerRoundness(
+                    mSharedState.showCornerRadiusInDesktopMode);
         } else {
-            mCornerRoundness.updateValue(TaskbarBackgroundRenderer.MAX_ROUNDNESS);
+            mCornerRoundness.value = TaskbarBackgroundRenderer.MAX_ROUNDNESS;
         }
+        updateCornerRoundness();
         onPostInit();
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
index a4fbb25..707d4b3 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
@@ -16,6 +16,7 @@
 package com.android.launcher3.taskbar;
 
 import static com.android.app.animation.Interpolators.EMPHASIZED;
+import static com.android.launcher3.Flags.enableScalingRevealHomeAnimation;
 import static com.android.launcher3.Hotseat.ALPHA_CHANNEL_TASKBAR_ALIGNMENT;
 import static com.android.launcher3.Hotseat.ALPHA_CHANNEL_TASKBAR_STASH;
 import static com.android.launcher3.LauncherState.HOTSEAT_ICONS;
@@ -42,6 +43,7 @@
 import android.animation.ObjectAnimator;
 import android.os.SystemClock;
 import android.util.Log;
+import android.view.animation.Interpolator;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -64,6 +66,7 @@
 import com.android.launcher3.util.MultiPropertyFactory.MultiProperty;
 import com.android.quickstep.RecentsAnimationCallbacks;
 import com.android.quickstep.RecentsAnimationController;
+import com.android.quickstep.util.ScalingWorkspaceRevealAnim;
 import com.android.quickstep.util.SystemUiFlagUtils;
 import com.android.quickstep.views.RecentsView;
 import com.android.systemui.animation.ViewRootSync;
@@ -682,7 +685,9 @@
             animatorSet.play(iconAlignAnim);
         }
 
-        animatorSet.setInterpolator(EMPHASIZED);
+        Interpolator interpolator = enableScalingRevealHomeAnimation()
+                ? ScalingWorkspaceRevealAnim.SCALE_INTERPOLATOR : EMPHASIZED;
+        animatorSet.setInterpolator(interpolator);
 
         if (start) {
             animatorSet.start();
@@ -781,6 +786,9 @@
     }
 
     protected void stashHotseat(boolean stash) {
+        // align taskbar with the hotseat icons before performing any animation
+        mControllers.taskbarViewController.setLauncherIconAlignment(/* alignmentRatio = */ 1,
+                mLauncher.getDeviceProfile());
         TaskbarStashController stashController = mControllers.taskbarStashController;
         stashController.updateStateForFlag(FLAG_STASHED_FOR_BUBBLES, stash);
         Runnable swapHotseatWithTaskbar = new Runnable() {
@@ -893,19 +901,6 @@
             mHotseatTranslationXAnimation.cancel();
             mHotseatTranslationXAnimation = null;
         }
-        Runnable postAnimationAction = new Runnable() {
-            @Override
-            public void run() {
-                mHotseatTranslationXAnimation = null;
-                // We only need to align the task bar when on launcher home screen
-                if (mControllers.taskbarStashController.isOnHome()) {
-                    mControllers.taskbarViewController.setLauncherIconAlignment(
-                            /* alignmentRatio = */ 1,
-                            mLauncher.getDeviceProfile()
-                    );
-                }
-            }
-        };
         Hotseat hotseat = mLauncher.getHotseat();
         AnimatorSet translationXAnimation = new AnimatorSet();
         MultiProperty iconsTranslationX = mLauncher.getHotseat()
@@ -928,14 +923,12 @@
             }
         }
         if (!animate) {
-            postAnimationAction.run();
             return;
         }
         mHotseatTranslationXAnimation = translationXAnimation;
         translationXAnimation.setStartDelay(FADE_OUT_ANIM_POSITION_DURATION_MS);
         translationXAnimation.setDuration(FADE_IN_ANIM_ALPHA_DURATION_MS);
         translationXAnimation.setInterpolator(Interpolators.EMPHASIZED);
-        translationXAnimation.addListener(AnimatorListeners.forEndCallback(postAnimationAction));
         translationXAnimation.start();
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index 78e7b47..f2f6500 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -16,7 +16,6 @@
 package com.android.launcher3.taskbar;
 
 import static android.content.Context.RECEIVER_NOT_EXPORTED;
-import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR;
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR_PANEL;
 
@@ -220,7 +219,7 @@
             TaskbarNavButtonCallbacks navCallbacks,
             @NonNull DesktopVisibilityController desktopVisibilityController) {
         Display display =
-                context.getSystemService(DisplayManager.class).getDisplay(DEFAULT_DISPLAY);
+                context.getSystemService(DisplayManager.class).getDisplay(context.getDisplayId());
         mContext = context.createWindowContext(display,
                 ENABLE_TASKBAR_NAVBAR_UNIFICATION ? TYPE_NAVIGATION_BAR : TYPE_NAVIGATION_BAR_PANEL,
                 null);
@@ -672,11 +671,6 @@
     @VisibleForTesting
     public void setSuspended(boolean isSuspended) {
         mIsSuspended = isSuspended;
-        if (mIsSuspended) {
-            removeTaskbarRootViewFromWindow();
-        } else {
-            addTaskbarRootViewToWindow();
-        }
     }
 
     private void addTaskbarRootViewToWindow() {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
index 57d4dbb..9c34ff0 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
@@ -16,7 +16,7 @@
 package com.android.launcher3.taskbar
 
 import android.content.Context
-import android.window.flags.DesktopModeFlags
+import android.window.DesktopModeFlags
 import androidx.annotation.VisibleForTesting
 import com.android.launcher3.Flags.enableRecentsInTaskbar
 import com.android.launcher3.model.data.ItemInfo
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index 8763509..ec94160 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -192,13 +192,44 @@
     }
 
     /**
-     // @return the maximum number of 'icons' that can fit in the taskbar.
-     // TODO(368119679): Assumes that they are all the same size.
+     * @return the maximum number of 'icons' that can fit in the taskbar.
      */
     private int calculateMaxNumIcons() {
-        int availableWidth = mActivityContext.getDeviceProfile().widthPx
-                - (mActivityContext.getDeviceProfile().edgeMarginPx * 2);
-        return Math.floorDiv(availableWidth, mIconTouchSize);
+        DeviceProfile deviceProfile = mActivityContext.getDeviceProfile();
+        int availableWidth = deviceProfile.widthPx;
+
+        // Reserve space required for edge margins, or for navbar if shown. If task bar needs to be
+        // center aligned with nav bar shown, reserve space on both sides.
+        availableWidth -= Math.max(deviceProfile.edgeMarginPx, deviceProfile.hotseatBarEndOffset);
+        availableWidth -= Math.max(deviceProfile.edgeMarginPx,
+                mShouldTryStartAlign ? 0 : deviceProfile.hotseatBarEndOffset);
+
+        // The space taken by an item icon used during layout.
+        int iconSize = 2 * mItemMarginLeftRight + mIconTouchSize;
+
+        int additionalIcons = 0;
+
+        if (mTaskbarDividerContainer != null) {
+            // Space for divider icon is reduced during layout compared to normal icon size, reserve
+            // space for the divider separately.
+            availableWidth -= iconSize - 4 * mItemMarginLeftRight;
+            ++additionalIcons;
+        }
+
+        // All apps icon takes less space compared to normal icon size, reserve space for the icon
+        // separately.
+        if (mAllAppsButtonContainer != null) {
+            boolean forceTransientTaskbarSize =
+                    enableTaskbarPinning() && !mActivityContext.isThreeButtonNav();
+            availableWidth -= iconSize - (int) getResources().getDimension(
+                    mAllAppsButtonContainer.getAllAppsButtonTranslationXOffset(
+                            forceTransientTaskbarSize || (
+                                    DisplayController.isTransientTaskbar(mActivityContext)
+                                            && !mActivityContext.isPhoneMode())));
+            ++additionalIcons;
+        }
+
+        return Math.floorDiv(availableWidth, iconSize) + additionalIcons;
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarItem.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarItem.kt
index 7a32ef1..680ffca 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarItem.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarItem.kt
@@ -17,10 +17,11 @@
 
 import android.graphics.Bitmap
 import android.graphics.Path
+import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutMessage
 import com.android.wm.shell.shared.bubbles.BubbleInfo
 
 /** An entity in the bubble bar. */
-sealed class BubbleBarItem(open var key: String, open var view: BubbleView)
+sealed class BubbleBarItem(open val key: String, open var view: BubbleView)
 
 /** Contains state info about a bubble in the bubble bar as well as presentation information. */
 data class BubbleBarBubble(
@@ -30,7 +31,8 @@
     var icon: Bitmap,
     var dotColor: Int,
     var dotPath: Path,
-    var appName: String
+    var appName: String,
+    var flyoutMessage: BubbleBarFlyoutMessage?,
 ) : BubbleBarItem(info.key, view)
 
 /** Represents the overflow bubble in the bubble bar. */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarSwipeController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarSwipeController.kt
index bc562a6..2d3642b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarSwipeController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarSwipeController.kt
@@ -21,15 +21,14 @@
 import androidx.annotation.VisibleForTesting
 import androidx.core.animation.doOnEnd
 import androidx.dynamicanimation.animation.SpringForce
-import com.android.launcher3.R
 import com.android.launcher3.anim.AnimatedFloat
 import com.android.launcher3.anim.SpringAnimationBuilder
 import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.taskbar.TaskbarThresholdUtils
-import com.android.launcher3.taskbar.bubbles.BubbleBarSwipeController.StartState.COLLAPSED
-import com.android.launcher3.taskbar.bubbles.BubbleBarSwipeController.StartState.EXPANDED
-import com.android.launcher3.taskbar.bubbles.BubbleBarSwipeController.StartState.STASHED
-import com.android.launcher3.taskbar.bubbles.BubbleBarSwipeController.StartState.UNKNOWN
+import com.android.launcher3.taskbar.bubbles.BubbleBarSwipeController.BarState.COLLAPSED
+import com.android.launcher3.taskbar.bubbles.BubbleBarSwipeController.BarState.EXPANDED
+import com.android.launcher3.taskbar.bubbles.BubbleBarSwipeController.BarState.STASHED
+import com.android.launcher3.taskbar.bubbles.BubbleBarSwipeController.BarState.UNKNOWN
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController
 import com.android.launcher3.touch.OverScroll
 
@@ -46,11 +45,9 @@
     private val animatedSwipeTranslation = AnimatedFloat(this::onSwipeUpdate)
 
     private val unstashThreshold: Int
-    private val expandThreshold: Int
     private val maxOverscroll: Int
-    private val stashThreshold: Int
 
-    private var swipeState: SwipeState = SwipeState()
+    private var swipeState: SwipeState = SwipeState(startState = UNKNOWN)
 
     constructor(tac: TaskbarActivityContext) : this(tac, DefaultDimensionProvider(tac))
 
@@ -58,9 +55,7 @@
     constructor(context: Context, dimensionProvider: DimensionProvider) {
         this.context = context
         unstashThreshold = dimensionProvider.unstashThreshold
-        expandThreshold = dimensionProvider.expandThreshold
         maxOverscroll = dimensionProvider.maxOverscroll
-        stashThreshold = dimensionProvider.stashThreshold
     }
 
     fun init(bubbleControllers: BubbleControllers) {
@@ -80,7 +75,7 @@
                 bubbleStashController.isBubbleBarVisible() -> COLLAPSED
                 else -> UNKNOWN
             }
-        swipeState = SwipeState(startState = startState)
+        swipeState = SwipeState(startState = startState, currentState = startState)
     }
 
     /** Update swipe distance to [dy] */
@@ -90,46 +85,25 @@
         }
         animatedSwipeTranslation.updateValue(dy)
 
-        val prevState = swipeState
-        // We can pass unstash threshold once per gesture, keep it true if it happened once
-        val passedUnstashThreshold = isUnstash(dy) || prevState.passedUnstashThreshold
-        // Expand happens at the end of the gesture, always keep the current value
-        val passedExpandThreshold = isExpand(dy)
-        // Stash happens at the end of the gesture, always keep the current value
-        val passedStashThreshold = isStash(dy)
-
-        if (
-            passedUnstashThreshold != prevState.passedUnstashThreshold ||
-                passedExpandThreshold != prevState.passedExpandThreshold ||
-                passedStashThreshold != prevState.passedStashThreshold
-        ) {
-            swipeState =
-                swipeState.copy(
-                    passedUnstashThreshold = passedUnstashThreshold,
-                    passedExpandThreshold = passedExpandThreshold,
-                    passedStashThreshold = passedStashThreshold,
-                )
-        }
-
-        if (
-            swipeState.startState == STASHED &&
-                swipeState.passedUnstashThreshold &&
-                !prevState.passedUnstashThreshold
-        ) {
-            bubbleStashController.showBubbleBar(expandBubbles = false)
+        swipeState.passedUnstash = isUnstash(dy)
+        // Tracking swipe gesture if we pass unstash threshold at least once during gesture
+        swipeState.isSwipe = swipeState.isSwipe || swipeState.passedUnstash
+        when {
+            canUnstash() && swipeState.passedUnstash -> {
+                swipeState.currentState = COLLAPSED
+                bubbleStashController.showBubbleBar(expandBubbles = false)
+            }
+            canStash() && !swipeState.passedUnstash -> {
+                swipeState.currentState = STASHED
+                bubbleStashController.stashBubbleBar()
+            }
         }
     }
 
     /** Finish tracking swipe gesture. Animate views back to resting state */
     fun finish() {
-        when {
-            swipeState.passedExpandThreshold &&
-                swipeState.startState in setOf(STASHED, COLLAPSED) -> {
-                bubbleStashController.showBubbleBar(expandBubbles = true)
-            }
-            swipeState.passedStashThreshold && swipeState.startState == COLLAPSED -> {
-                bubbleStashController.stashBubbleBar()
-            }
+        if (swipeState.passedUnstash && swipeState.startState in setOf(STASHED, COLLAPSED)) {
+            bubbleStashController.showBubbleBar(expandBubbles = true)
         }
         if (animatedSwipeTranslation.value == 0f) {
             reset()
@@ -140,15 +114,21 @@
 
     /** Returns `true` if we are tracking a swipe gesture */
     fun isSwipeGesture(): Boolean {
-        return swipeState.passedUnstashThreshold ||
-            swipeState.passedExpandThreshold ||
-            swipeState.passedStashThreshold
+        return swipeState.isSwipe
     }
 
     private fun canHandleSwipe(dy: Float): Boolean {
         return when (swipeState.startState) {
-            STASHED -> dy < 0 // stashed bar only handles swipe up
-            COLLAPSED -> true // collapsed bar can be swiped in either direction
+            STASHED -> {
+                if (swipeState.currentState == COLLAPSED) {
+                    // if we have unstashed the bar, allow swipe in both directions
+                    true
+                } else {
+                    // otherwise, only allow swipe up on stash handle
+                    dy < 0
+                }
+            }
+            COLLAPSED -> dy < 0 // collapsed bar can only be swiped up
             UNKNOWN,
             EXPANDED -> false // expanded bar can't be swiped
         }
@@ -158,12 +138,13 @@
         return dy < -unstashThreshold
     }
 
-    private fun isExpand(dy: Float): Boolean {
-        return dy < -expandThreshold
+    private fun canStash(): Boolean {
+        // Only allow stashing if we started from stashed state
+        return swipeState.startState == STASHED && swipeState.currentState == COLLAPSED
     }
 
-    private fun isStash(dy: Float): Boolean {
-        return dy > stashThreshold
+    private fun canUnstash(): Boolean {
+        return swipeState.currentState == STASHED
     }
 
     private fun reset() {
@@ -175,7 +156,7 @@
             }
         }
         springAnimation = null
-        swipeState = SwipeState()
+        swipeState = SwipeState(startState = UNKNOWN)
     }
 
     private fun onSwipeUpdate(value: Float) {
@@ -197,13 +178,13 @@
     }
 
     internal data class SwipeState(
-        val startState: StartState = UNKNOWN,
-        val passedUnstashThreshold: Boolean = false,
-        val passedExpandThreshold: Boolean = false,
-        val passedStashThreshold: Boolean = false,
+        val startState: BarState,
+        var currentState: BarState = UNKNOWN,
+        var passedUnstash: Boolean = false,
+        var isSwipe: Boolean = false,
     )
 
-    internal enum class StartState {
+    internal enum class BarState {
         UNKNOWN,
         STASHED,
         COLLAPSED,
@@ -214,17 +195,13 @@
     @VisibleForTesting
     interface DimensionProvider {
         val unstashThreshold: Int
-        val expandThreshold: Int
         val maxOverscroll: Int
-        val stashThreshold: Int
     }
 
     private class DefaultDimensionProvider(taskbarActivityContext: TaskbarActivityContext) :
         DimensionProvider {
         override val unstashThreshold: Int
-        override val expandThreshold: Int
         override val maxOverscroll: Int
-        override val stashThreshold: Int
 
         init {
             val resources = taskbarActivityContext.resources
@@ -233,14 +210,7 @@
                     resources,
                     taskbarActivityContext.deviceProfile,
                 )
-            // TODO(325673340): review threshold with ux
-            expandThreshold =
-                TaskbarThresholdUtils.getAppWindowThreshold(
-                    resources,
-                    taskbarActivityContext.deviceProfile,
-                )
             maxOverscroll = taskbarActivityContext.deviceProfile.heightPx - unstashThreshold
-            stashThreshold = resources.getDimensionPixelSize(R.dimen.taskbar_to_nav_threshold)
         }
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index 7fed381..c5d649e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -1492,6 +1492,38 @@
         return bubbles;
     }
 
+    /**
+     * Returns the distance between the top left corner of the bubble bar to the center of the dot
+     * of the selected bubble.
+     */
+    PointF getSelectedBubbleDotDistanceFromTopLeft() {
+        if (mSelectedBubbleView == null) {
+            return new PointF(0, 0);
+        }
+        final int indexOfSelectedBubble = indexOfChild(mSelectedBubbleView);
+        final boolean onLeft = mBubbleBarLocation.isOnLeft(isLayoutRtl());
+        final float selectedBubbleTx = isExpanded()
+                ? getExpandedBubbleTranslationX(indexOfSelectedBubble, getChildCount(), onLeft)
+                : getCollapsedBubbleTranslationX(indexOfSelectedBubble, getChildCount(), onLeft);
+        PointF selectedBubbleDotCenter = mSelectedBubbleView.getDotCenter();
+
+        return new PointF(
+                selectedBubbleTx + selectedBubbleDotCenter.x,
+                mBubbleBarPadding + mPointerSize + selectedBubbleDotCenter.y);
+    }
+
+    int getSelectedBubbleDotColor() {
+        return mSelectedBubbleView == null ? 0 : mSelectedBubbleView.getDotColor();
+    }
+
+    int getPointerSize() {
+        return mPointerSize;
+    }
+
+    float getBubbleElevation() {
+        return mBubbleElevation;
+    }
+
     /** Interface for BubbleBarView to communicate with its controller. */
     interface Controller {
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index c164eec..69e1d43 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -29,6 +29,7 @@
 import android.util.TypedValue;
 import android.view.MotionEvent;
 import android.view.View;
+import android.widget.FrameLayout;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -42,6 +43,7 @@
 import com.android.launcher3.taskbar.TaskbarInsetsController;
 import com.android.launcher3.taskbar.TaskbarStashController;
 import com.android.launcher3.taskbar.bubbles.animation.BubbleBarViewAnimator;
+import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutPositioner;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController;
 import com.android.launcher3.util.MultiPropertyFactory;
 import com.android.launcher3.util.MultiValueAlpha;
@@ -63,6 +65,8 @@
     private static final float APP_ICON_SMALL_DP = 44f;
     private static final float APP_ICON_MEDIUM_DP = 48f;
     private static final float APP_ICON_LARGE_DP = 52f;
+    /** The dot size is defined as a percentage of the icon size. */
+    private static final float DOT_TO_BUBBLE_SIZE_RATIO = 0.228f;
     private final SystemUiProxy mSystemUiProxy;
     private final TaskbarActivityContext mActivity;
     private final BubbleBarView mBarView;
@@ -119,7 +123,8 @@
     @Nullable
     private BubbleBarBoundsChangeListener mBoundsChangeListener;
 
-    public BubbleBarViewController(TaskbarActivityContext activity, BubbleBarView barView) {
+    public BubbleBarViewController(TaskbarActivityContext activity, BubbleBarView barView,
+            FrameLayout bubbleBarContainer) {
         mActivity = activity;
         mBarView = barView;
         mSystemUiProxy = SystemUiProxy.INSTANCE.get(mActivity);
@@ -208,6 +213,59 @@
         };
     }
 
+    private BubbleBarFlyoutPositioner createFlyoutPositioner() {
+        return new BubbleBarFlyoutPositioner() {
+
+            @Override
+            public boolean isOnLeft() {
+                return mBarView.getBubbleBarLocation().isOnLeft(mBarView.isLayoutRtl());
+            }
+
+            @Override
+            public float getTargetTy() {
+                return mBarView.getTranslationY() - mBarView.getHeight();
+            }
+
+            @Override
+            @NonNull
+            public PointF getDistanceToCollapsedPosition() {
+                // the flyout animates from the selected bubble dot. calculate the distance it needs
+                // to translate itself to its starting position.
+                PointF distanceToDotCenter = mBarView.getSelectedBubbleDotDistanceFromTopLeft();
+
+                // if we're gravitating left, return the distance between the top left corner of the
+                // bubble bar and the bottom left corner of the dot.
+                // if we're gravitating right, return the distance between the top right corner of
+                // the bubble bar and the bottom right corner of the dot.
+                float distanceX = isOnLeft()
+                        ? distanceToDotCenter.x - getCollapsedSize() / 2
+                        : mBarView.getWidth() - distanceToDotCenter.x - getCollapsedSize() / 2;
+                float distanceY = distanceToDotCenter.y + getCollapsedSize() / 2;
+                return new PointF(distanceX, distanceY);
+            }
+
+            @Override
+            public float getCollapsedSize() {
+                return mIconSize * DOT_TO_BUBBLE_SIZE_RATIO;
+            }
+
+            @Override
+            public int getCollapsedColor() {
+                return mBarView.getSelectedBubbleDotColor();
+            }
+
+            @Override
+            public float getCollapsedElevation() {
+                return mBarView.getBubbleElevation();
+            }
+
+            @Override
+            public float getDistanceToRevealTriangle() {
+                return getDistanceToCollapsedPosition().y - mBarView.getPointerSize();
+            }
+        };
+    }
+
     private void onBubbleClicked(BubbleView bubbleView) {
         bubbleView.markSeen();
         BubbleBarItem bubble = bubbleView.getBubble();
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleCreator.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleCreator.java
index 340a120..c5efe2f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleCreator.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleCreator.java
@@ -22,6 +22,7 @@
 import static android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_PINNED_BY_ANY_LAUNCHER;
 
 import static com.android.launcher3.icons.FastBitmapDrawable.WHITE_SCRIM_ALPHA;
+import static com.android.wm.shell.shared.bubbles.FlyoutDrawableLoader.loadFlyoutDrawable;
 
 import android.annotation.Nullable;
 import android.content.Context;
@@ -49,7 +50,9 @@
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.BubbleIconFactory;
 import com.android.launcher3.shortcuts.ShortcutRequest;
+import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutMessage;
 import com.android.wm.shell.shared.bubbles.BubbleInfo;
+import com.android.wm.shell.shared.bubbles.ParcelableFlyoutMessage;
 
 /**
  * Loads the necessary info to populate / present a bubble (name, icon, shortcut).
@@ -157,13 +160,16 @@
         dotColor = ColorUtils.blendARGB(badgeBitmapInfo.color,
                 Color.WHITE, WHITE_SCRIM_ALPHA / 255f);
 
+        final BubbleBarFlyoutMessage flyoutMessage =
+                getFlyoutMessage(info.getParcelableFlyoutMessage());
+
         if (existingBubble == null) {
             LayoutInflater inflater = LayoutInflater.from(context);
             BubbleView bubbleView = (BubbleView) inflater.inflate(
                     R.layout.bubblebar_item_view, barView, false /* attachToRoot */);
 
             BubbleBarBubble bubble = new BubbleBarBubble(info, bubbleView,
-                    badgeBitmap, bubbleBitmap, dotColor, dotPath, appName);
+                    badgeBitmap, bubbleBitmap, dotColor, dotPath, appName, flyoutMessage);
             bubbleView.setBubble(bubble);
             return bubble;
         } else {
@@ -174,10 +180,25 @@
             existingBubble.setDotColor(dotColor);
             existingBubble.setDotPath(dotPath);
             existingBubble.setAppName(appName);
+            existingBubble.setFlyoutMessage(flyoutMessage);
             return existingBubble;
         }
     }
 
+    @Nullable
+    private BubbleBarFlyoutMessage getFlyoutMessage(
+            @Nullable ParcelableFlyoutMessage parcelableFlyoutMessage) {
+        if (parcelableFlyoutMessage == null) {
+            return null;
+        }
+        String title = parcelableFlyoutMessage.getTitle();
+        String message = parcelableFlyoutMessage.getMessage();
+        return new BubbleBarFlyoutMessage(
+                loadFlyoutDrawable(parcelableFlyoutMessage.getIcon(), mContext),
+                title == null ? "" : title,
+                message == null ? "" : message);
+    }
+
     /**
      * Creates the overflow view shown in the bubble bar.
      *
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
index 561df5c..707655c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
@@ -22,6 +22,7 @@
 import android.graphics.Canvas;
 import android.graphics.Color;
 import android.graphics.Path;
+import android.graphics.PointF;
 import android.graphics.Rect;
 import android.graphics.drawable.BitmapDrawable;
 import android.os.Bundle;
@@ -67,8 +68,7 @@
     private float mAnimatingToDotScale;
     // The current scale value of the dot
     private float mDotScale;
-
-    private boolean mProvideShadowOutline = true;
+    private boolean mDotSuppressedForBubbleUpdate = false;
 
     // TODO: (b/273310265) handle RTL
     // Whether the bubbles are positioned on the left or right side of the screen
@@ -300,6 +300,10 @@
     }
 
     void updateDotVisibility(boolean animate) {
+        if (mDotSuppressedForBubbleUpdate) {
+            // if the dot is suppressed for
+            return;
+        }
         final float targetScale = hasUnseenContent() ? 1f : 0f;
         if (animate) {
             animateDotScale(targetScale);
@@ -317,6 +321,20 @@
         }
     }
 
+    /**
+     * Suppresses or un-suppresses drawing the dot due to an update for this bubble.
+     *
+     * <p>If the dot is being suppressed and is already visible, it remains visible because it is
+     * used as a starting point for the animation. If the dot is being unsuppressed, it is
+     * redrawn if needed.
+     */
+    public void suppressDotForBubbleUpdate(boolean suppress) {
+        mDotSuppressedForBubbleUpdate = suppress;
+        if (!suppress) {
+            showDotIfNeeded(/* animate= */ false);
+        }
+    }
+
     boolean hasUnseenContent() {
         return mBubble != null
                 && mBubble instanceof BubbleBarBubble
@@ -353,8 +371,8 @@
     }
 
     void showDotIfNeeded(boolean animate) {
-        // only show the dot if we have unseen content
-        if (!hasUnseenContent()) {
+        // only show the dot if we have unseen content and it's not suppressed
+        if (!hasUnseenContent() || mDotSuppressedForBubbleUpdate) {
             return;
         }
         if (animate) {
@@ -406,6 +424,23 @@
                 }).start();
     }
 
+    /**
+     * Returns the distance from the top left corner of this bubble view to the center of its dot.
+     */
+    public PointF getDotCenter() {
+        float[] dotPosition =
+                mOnLeft ? mDotRenderer.getLeftDotPosition() : mDotRenderer.getRightDotPosition();
+        getDrawingRect(mTempBounds);
+        float dotCenterX = mTempBounds.width() * dotPosition[0];
+        float dotCenterY = mTempBounds.height() * dotPosition[1];
+        return new PointF(dotCenterX, dotCenterY);
+    }
+
+    /** Returns the dot color. */
+    public int getDotColor() {
+        return mDotColor;
+    }
+
     @Override
     public String toString() {
         String toString = mBubble != null ? mBubble.getKey() : "null";
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 0e7abf2..39bf6ac 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -19,7 +19,7 @@
 import static android.os.Trace.TRACE_TAG_APP;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_OPTIMIZE_MEASURE;
 import static android.view.accessibility.AccessibilityEvent.TYPE_VIEW_FOCUSED;
-import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY;
+import static android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY;
 
 import static com.android.app.animation.Interpolators.EMPHASIZED;
 import static com.android.internal.jank.Cuj.CUJ_LAUNCHER_LAUNCH_APP_PAIR_FROM_WORKSPACE;
@@ -1316,6 +1316,10 @@
         mTISBindHelper.setPredictiveBackToHomeInProgress(isInProgress);
     }
 
+    public boolean getPredictiveBackToHomeInProgress() {
+        return mIsPredictiveBackToHomeInProgress;
+    }
+
     @Override
     public boolean areDesktopTasksVisible() {
         DesktopVisibilityController desktopVisibilityController = getDesktopVisibilityController();
diff --git a/quickstep/src/com/android/launcher3/uioverrides/flags/DevOptionsUiHelper.kt b/quickstep/src/com/android/launcher3/uioverrides/flags/DevOptionsUiHelper.kt
index 181cba0..417bb74 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/flags/DevOptionsUiHelper.kt
+++ b/quickstep/src/com/android/launcher3/uioverrides/flags/DevOptionsUiHelper.kt
@@ -35,9 +35,6 @@
 import android.provider.Settings.Secure
 import android.text.Html
 import android.util.AttributeSet
-import android.util.Base64
-import android.util.Base64.NO_PADDING
-import android.util.Base64.NO_WRAP
 import android.view.inputmethod.EditorInfo
 import android.widget.TextView
 import android.widget.Toast
@@ -57,9 +54,10 @@
 import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET
 import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT
 import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
-import com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_KEY
 import com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_LABEL
 import com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_TAG
+import com.android.launcher3.LauncherSettings.Settings.LAYOUT_PROVIDER_KEY
+import com.android.launcher3.LauncherSettings.Settings.createBlobProviderKey
 import com.android.launcher3.R
 import com.android.launcher3.model.data.FolderInfo
 import com.android.launcher3.model.data.ItemInfo
@@ -241,7 +239,7 @@
     private fun DebugInfo<Boolean>.getBoolValue() =
         DeviceConfigHelper.prefs.getBoolean(
             this.key,
-            DeviceConfig.getBoolean(NAMESPACE_LAUNCHER, this.key, this.valueInCode)
+            DeviceConfig.getBoolean(NAMESPACE_LAUNCHER, this.key, this.valueInCode),
         )
 
     private fun DebugInfo<Int>.getIntValueAsString() =
@@ -265,7 +263,7 @@
         val pluginPermissionApps =
             pm.getPackagesHoldingPermissions(
                     arrayOf(PLUGIN_PERMISSION),
-                    PackageManager.MATCH_DISABLED_COMPONENTS
+                    PackageManager.MATCH_DISABLED_COMPONENTS,
                 )
                 .map { it.packageName }
 
@@ -274,7 +272,7 @@
                 pm.queryIntentServices(
                         Intent(action),
                         PackageManager.MATCH_DISABLED_COMPONENTS or
-                            PackageManager.GET_RESOLVED_FILTER
+                            PackageManager.GET_RESOLVED_FILTER,
                     )
                     .filter { pluginPermissionApps.contains(it.serviceInfo.packageName) }
             }
@@ -316,7 +314,7 @@
                             infoList.forEach {
                                 manager.pluginEnabler.setDisabled(
                                     it.serviceInfo.componentName,
-                                    disabledState
+                                    disabledState,
                                 )
                             }
                             manager.notifyChange(Intent(Intent.ACTION_PACKAGE_CHANGED, pluginUri))
@@ -387,12 +385,12 @@
             addOnboardPref(
                 "All Apps Bounce",
                 HOME_BOUNCE_SEEN.sharedPrefKey,
-                HOME_BOUNCE_COUNT.sharedPrefKey
+                HOME_BOUNCE_COUNT.sharedPrefKey,
             )
             addOnboardPref(
                 "Hybrid Hotseat Education",
                 HOTSEAT_DISCOVERY_TIP_COUNT.sharedPrefKey,
-                HOTSEAT_LONGPRESS_TIP_SEEN.sharedPrefKey
+                HOTSEAT_LONGPRESS_TIP_SEEN.sharedPrefKey,
             )
             addOnboardPref("Taskbar Education", TASKBAR_EDU_TOOLTIP_STEP.sharedPrefKey)
             addOnboardPref("Taskbar Search Education", TASKBAR_SEARCH_EDU_SEEN.sharedPrefKey)
@@ -470,13 +468,16 @@
                         session.allowPublicAccess()
 
                         session.commit(ORDERED_BG_EXECUTOR) {
-                            val key = Base64.encodeToString(digest, NO_WRAP or NO_PADDING)
-                            Secure.putString(resolver, LAYOUT_DIGEST_KEY, key)
+                            Secure.putString(
+                                resolver,
+                                LAYOUT_PROVIDER_KEY,
+                                createBlobProviderKey(digest),
+                            )
 
                             MODEL_EXECUTOR.submit { model.modelDbController.createEmptyDB() }.get()
                             MAIN_EXECUTOR.submit { model.forceReload() }.get()
                             MODEL_EXECUTOR.submit {}.get()
-                            Secure.putString(resolver, LAYOUT_DIGEST_KEY, null)
+                            Secure.putString(resolver, LAYOUT_PROVIDER_KEY, null)
                         }
                     }
                 }
@@ -512,7 +513,7 @@
                     info.providerName.className,
                     info.spanX,
                     info.spanY,
-                    userType
+                    userType,
                 )
         }
     }
@@ -520,7 +521,7 @@
     private fun createUriPickerIntent(
         action: String,
         executor: Executor,
-        callback: (uri: Uri) -> Unit
+        callback: (uri: Uri) -> Unit,
     ): Intent {
         val pendingIntent =
             PendingIntent(
@@ -532,7 +533,7 @@
                         allowlistToken: IBinder?,
                         finishedReceiver: IIntentReceiver?,
                         requiredPermission: String?,
-                        options: Bundle?
+                        options: Bundle?,
                     ) {
                         intent.data?.let { uri -> executor.execute { callback(uri) } }
                     }
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 3413532..fbb2c06 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -88,8 +88,8 @@
 import android.view.WindowInsets;
 import android.view.animation.Interpolator;
 import android.widget.Toast;
+import android.window.DesktopModeFlags;
 import android.window.PictureInPictureSurfaceTransaction;
-import android.window.flags.DesktopModeFlags;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -156,6 +156,8 @@
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.shared.startingsurface.SplashScreenExitAnimationUtils;
 
+import kotlin.Unit;
+
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -165,8 +167,6 @@
 import java.util.OptionalInt;
 import java.util.function.Consumer;
 
-import kotlin.Unit;
-
 /**
  * Handles the navigation gestures when Launcher is the default home activity.
  */
diff --git a/quickstep/src/com/android/quickstep/BaseContainerInterface.java b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
index e625f55..143ef12 100644
--- a/quickstep/src/com/android/quickstep/BaseContainerInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
@@ -46,7 +46,6 @@
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.views.ScrimView;
-import com.android.quickstep.fallback.window.RecentsWindowManager;
 import com.android.quickstep.orientation.RecentsPagedOrientationHandler;
 import com.android.quickstep.util.ActivityInitListener;
 import com.android.quickstep.util.AnimatorControllerWithResistance;
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index a55cf18..7a0d701 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -46,11 +46,11 @@
 import android.view.MotionEvent;
 import android.view.RemoteAnimationTarget;
 import android.view.SurfaceControl;
+import android.window.DesktopModeFlags;
 import android.window.IOnBackInvokedCallback;
 import android.window.RemoteTransition;
 import android.window.TaskSnapshot;
 import android.window.TransitionFilter;
-import android.window.flags.DesktopModeFlags;
 
 import androidx.annotation.MainThread;
 import androidx.annotation.Nullable;
@@ -1079,16 +1079,6 @@
         }
     }
 
-    public void removeFromSideStage(int taskId) {
-        if (mSplitScreen != null) {
-            try {
-                mSplitScreen.removeFromSideStage(taskId);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call removeFromSideStage");
-            }
-        }
-    }
-
     //
     // One handed
     //
@@ -1492,6 +1482,17 @@
         }
     }
 
+    /** Call shell to remove the desktop that is on given `displayId` */
+    public void removeDesktop(int displayId) {
+        if (mDesktopMode != null) {
+            try {
+                mDesktopMode.removeDesktop(displayId);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call removeDesktop", e);
+            }
+        }
+    }
+
     //
     // Unfold transition
     //
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index 98d7628..bda292a 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -298,7 +298,7 @@
                 && Flags.enableFallbackOverviewInWindow()){
             mRecentsAnimationStartPending =
                     getSystemUiProxy().startRecentsActivity(intent, options, mCallbacks);
-            mRecentsWindowsManager.startRecentsWindow();
+            mRecentsWindowsManager.startRecentsWindow(mCallbacks);
         } else {
             options.setPendingIntentBackgroundActivityStartMode(
                     ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS);
@@ -485,10 +485,6 @@
         mTargets = null;
         mLastGestureState = null;
         mLastAppearedTaskTargets = null;
-
-        if(Flags.enableFallbackOverviewInWindow()) {
-            mRecentsWindowsManager.cleanup();
-        }
     }
 
     @Nullable
diff --git a/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java b/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java
index eef1f96..341c868 100644
--- a/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java
+++ b/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java
@@ -20,6 +20,7 @@
 import com.android.launcher3.dagger.LauncherBaseAppComponent;
 import com.android.launcher3.model.WellbeingModel;
 import com.android.quickstep.logging.SettingsChangeLogger;
+import com.android.quickstep.util.AsyncClockEventDelegate;
 
 /**
  * Launcher Quickstep base component for Dagger injection.
@@ -33,4 +34,6 @@
     SettingsChangeLogger getSettingsChangeLogger();
 
     WellbeingModel getWellbeingModel();
+
+    AsyncClockEventDelegate getAsyncClockEventDelegate();
 }
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
index b5830fd..fbf671f 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
@@ -46,6 +46,9 @@
 import com.android.launcher3.views.ScrimView
 import com.android.quickstep.FallbackWindowInterface
 import com.android.quickstep.OverviewComponentObserver
+import com.android.quickstep.RecentsAnimationCallbacks
+import com.android.quickstep.RecentsAnimationCallbacks.RecentsAnimationListener
+import com.android.quickstep.RecentsAnimationController
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.RemoteAnimationTargets
 import com.android.quickstep.SystemUiProxy
@@ -66,12 +69,16 @@
 import com.android.quickstep.views.OverviewActionsView
 import com.android.quickstep.views.RecentsView
 import com.android.quickstep.views.RecentsViewContainer
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.android.systemui.shared.system.TaskStackChangeListener
+import com.android.systemui.shared.system.TaskStackChangeListeners
 import java.util.function.Predicate
 
 /**
  * Class that will manage RecentsView lifecycle within a window and interface correctly where
- * needed. This allows us to run RecentsView in a window where needed. todo: b/365776320,
- * b/365777482
+ * needed. This allows us to run RecentsView in a window where needed.
+ *
+ * todo: b/365776320,b/365777482
  *
  * To add new protologs, see [RecentsWindowProtoLogProxy]. To enable logging to logcat, see
  * [QuickstepProtoLogGroup.Constants.DEBUG_RECENTS_WINDOW]
@@ -98,22 +105,48 @@
     private var actionsView: OverviewActionsView<*>? = null
     private var scrimView: ScrimView? = null
 
-    private var isShown = false
+    private var callbacks: RecentsAnimationCallbacks? = null
 
     private var tisBindHelper: TISBindHelper = TISBindHelper(this) {}
 
     // Callback array that corresponds to events defined in @ActivityEvent
     private val mEventCallbacks =
-        arrayOf(RunnableList(), RunnableList(), RunnableList(), RunnableList())
+        listOf(RunnableList(), RunnableList(), RunnableList(), RunnableList())
     private var onInitListener: Predicate<Boolean>? = null
 
+    private val taskStackChangeListener =
+        object : TaskStackChangeListener {
+            override fun onTaskMovedToFront(taskId: Int) {
+                if ((isShowing() && isInState(DEFAULT))) {
+                    // handling state where we end recents animation by swiping livetile away
+                    // TODO: animate this switch.
+                    cleanupRecentsWindow()
+                }
+            }
+        }
+
+    private val recentsAnimationListener =
+        object : RecentsAnimationListener {
+            override fun onRecentsAnimationCanceled(thumbnailDatas: HashMap<Int, ThumbnailData>) {
+                recentAnimationStopped()
+            }
+
+            override fun onRecentsAnimationFinished(controller: RecentsAnimationController) {
+                recentAnimationStopped()
+            }
+        }
+
     init {
         FallbackWindowInterface.init(this)
+        TaskStackChangeListeners.getInstance().registerTaskStackListener(taskStackChangeListener)
     }
 
     override fun destroy() {
         super.destroy()
+        cleanupRecentsWindow()
         FallbackWindowInterface.getInstance()?.destroy()
+        TaskStackChangeListeners.getInstance().unregisterTaskStackListener(taskStackChangeListener)
+        callbacks?.removeListener(recentsAnimationListener)
     }
 
     override fun startHome() {
@@ -135,6 +168,7 @@
                 ),
             )
         OverviewComponentObserver.startHomeIntentSafely(this, options.toBundle(), TAG)
+        stateManager.moveToRestState()
     }
 
     private val mAnimationToHomeFactory =
@@ -164,29 +198,35 @@
                 anim,
                 this@RecentsWindowManager,
                 {
-                    getStateManager().goToState(HOME, false)
-                    cleanup()
+                    getStateManager().goToState(BG_LAUNCHER, false)
+                    cleanupRecentsWindow()
                 },
                 true, /* skipFirstFrame */
             )
         }
 
-    fun cleanup() {
-        RecentsWindowProtoLogProxy.logCleanup(isShown)
-        if (isShown) {
+    private fun cleanupRecentsWindow() {
+        RecentsWindowProtoLogProxy.logCleanup(isShowing())
+        if (isShowing()) {
             windowManager.removeViewImmediate(windowView)
-            isShown = false
         }
+        stateManager.moveToRestState()
+        callbacks?.removeListener(recentsAnimationListener)
     }
 
-    fun startRecentsWindow() {
-        RecentsWindowProtoLogProxy.logStartRecentsWindow(isShown, windowView == null)
-        if (isShown) return
+    private fun isShowing(): Boolean {
+        return windowView?.parent != null
+    }
+
+    fun startRecentsWindow(callbacks: RecentsAnimationCallbacks? = null) {
+        RecentsWindowProtoLogProxy.logStartRecentsWindow(isShowing(), windowView == null)
+        if (isShowing()) {
+            return
+        }
         if (windowView == null) {
             windowView = layoutInflater.inflate(R.layout.fallback_recents_activity, null)
         }
         windowManager.addView(windowView, windowLayoutParams)
-        isShown = true
 
         windowView?.systemUiVisibility =
             (View.SYSTEM_UI_FLAG_LAYOUT_STABLE or
@@ -215,6 +255,15 @@
 
         mSystemUiController = SystemUiController(windowView)
         onInitListener?.test(true)
+
+        this.callbacks = callbacks
+        callbacks?.addListener(recentsAnimationListener)
+    }
+
+    private fun recentAnimationStopped() {
+        if (isInState(BACKGROUND_APP)) {
+            cleanupRecentsWindow()
+        }
     }
 
     override fun canStartHomeSafely(): Boolean {
@@ -246,14 +295,18 @@
         return stateManager.state == state
     }
 
-    override fun onStateSetStart(state: RecentsState?) {
+    override fun onStateSetStart(state: RecentsState) {
         super.onStateSetStart(state)
         RecentsWindowProtoLogProxy.logOnStateSetStart(getStateName(state))
     }
 
-    override fun onStateSetEnd(state: RecentsState?) {
+    override fun onStateSetEnd(state: RecentsState) {
         super.onStateSetEnd(state)
         RecentsWindowProtoLogProxy.logOnStateSetEnd(getStateName(state))
+
+        if (state == HOME || state == BG_LAUNCHER) {
+            cleanupRecentsWindow()
+        }
     }
 
     private fun getStateName(state: RecentsState?): String {
@@ -319,7 +372,7 @@
     }
 
     override fun isStarted(): Boolean {
-        return isShown
+        return isShowing() && isInState(DEFAULT)
     }
 
     /** Adds a callback for the provided activity event */
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index 1d4160d..2daaaf9 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -386,12 +386,12 @@
                 // and then write to StatsLog.
                 app.getModel().enqueueModelUpdateTask((taskController, dataModel, apps) ->
                         write(event, applyOverwrites(mItemInfo.buildProto(
-                                dataModel.collections.get(mItemInfo.container)))));
+                                dataModel.collections.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)
                 Executors.MODEL_EXECUTOR.execute(
-                        () -> write(event, applyOverwrites(mItemInfo.buildProto())));
+                        () -> write(event, applyOverwrites(mItemInfo.buildProto(mContext))));
             }
         }
 
diff --git a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
index b53650e..44b8b8d 100644
--- a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
+++ b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
@@ -34,6 +34,7 @@
 import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.quickstep.task.viewmodel.TaskOverlayViewModel
 import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
+import com.android.quickstep.task.viewmodel.TaskThumbnailViewModelImpl
 import com.android.quickstep.task.viewmodel.TaskViewData
 import com.android.quickstep.task.viewmodel.TaskViewModel
 import com.android.quickstep.views.TaskViewType
@@ -180,7 +181,7 @@
                 TaskContainerData::class.java -> TaskContainerData()
                 TaskThumbnailViewData::class.java -> TaskThumbnailViewData()
                 TaskThumbnailViewModel::class.java ->
-                    TaskThumbnailViewModel(
+                    TaskThumbnailViewModelImpl(
                         recentsViewData = inject(),
                         taskViewData = inject(scopeId, extras),
                         taskContainerData = inject(scopeId),
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
index 4970685..f55462a 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
@@ -10,144 +10,40 @@
  * 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 goveryning permissions and
+ * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 
 package com.android.quickstep.task.viewmodel
 
-import android.annotation.ColorInt
-import android.app.ActivityTaskManager.INVALID_TASK_ID
 import android.graphics.Matrix
-import android.util.Log
-import androidx.core.graphics.ColorUtils
-import com.android.quickstep.recents.data.RecentTasksRepository
-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.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.asStateFlow
-import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.flatMapLatest
-import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.flow.StateFlow
 
-@OptIn(ExperimentalCoroutinesApi::class)
-class TaskThumbnailViewModel(
-    recentsViewData: RecentsViewData,
-    taskViewData: TaskViewData,
-    taskContainerData: TaskContainerData,
-    private val tasksRepository: RecentTasksRepository,
-    private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase,
-    private val splashAlphaUseCase: SplashAlphaUseCase,
-) {
-    private val task = MutableStateFlow<Flow<Task?>>(flowOf(null))
-    private val splashProgress = MutableStateFlow(flowOf(0f))
-    private var taskId: Int = INVALID_TASK_ID
-
+/** ViewModel for representing TaskThumbnails */
+interface TaskThumbnailViewModel {
     /**
      * Progress for changes in corner radius. progress: 0 = overview corner radius; 1 = fullscreen
      * corner radius.
      */
-    val cornerRadiusProgress =
-        if (taskViewData.isOutlineFormedByThumbnailView) recentsViewData.fullscreenProgress
-        else MutableStateFlow(1f).asStateFlow()
+    val cornerRadiusProgress: StateFlow<Float>
 
-    val inheritedScale =
-        combine(recentsViewData.scale, taskViewData.scale) { recentsScale, taskScale ->
-            recentsScale * taskScale
-        }
+    /** The accumulated View.scale value for parent Views up to and including RecentsView */
+    val inheritedScale: Flow<Float>
 
-    val dimProgress: Flow<Float> =
-        combine(taskContainerData.taskMenuOpenProgress, recentsViewData.tintAmount) {
-            taskMenuOpenProgress,
-            tintAmount ->
-            max(taskMenuOpenProgress * MAX_SCRIM_ALPHA, tintAmount)
-        }
-    val splashAlpha = splashProgress.flatMapLatest { it }
+    /** Provides the level of dimming that the View should have */
+    val dimProgress: Flow<Float>
 
-    private val isLiveTile =
-        combine(
-                task.flatMapLatest { it }.map { it?.key?.id }.distinctUntilChanged(),
-                recentsViewData.runningTaskIds,
-                recentsViewData.runningTaskShowScreenshot,
-            ) { taskId, runningTaskIds, runningTaskShowScreenshot ->
-                runningTaskIds.contains(taskId) && !runningTaskShowScreenshot
-            }
-            .distinctUntilChanged()
+    /** Provides the alpha of the splash icon */
+    val splashAlpha: Flow<Float>
 
-    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 -> LiveTile
-                    isBackgroundOnly(taskVal) ->
-                        BackgroundOnly(taskVal.colorBackground.removeAlpha())
-                    isSnapshotSplashState(taskVal) ->
-                        SnapshotSplash(createSnapshotState(taskVal), taskVal.icon)
-                    else -> Uninitialized
-                }
-            }
-            .distinctUntilChanged()
+    /** Provides the UiState by which the task thumbnail can be represented */
+    val uiState: Flow<TaskThumbnailUiState>
 
-    fun bind(taskId: Int) {
-        Log.d(TAG, "bind taskId: $taskId")
-        this.taskId = taskId
-        task.value = tasksRepository.getTaskDataById(taskId)
-        splashProgress.value = splashAlphaUseCase.execute(taskId)
-    }
+    /** Attaches this ViewModel to a specific task id for it to provide data from. */
+    fun bind(taskId: Int)
 
-    fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix {
-        return runBlocking {
-            when (
-                val thumbnailPositionState =
-                    getThumbnailPositionUseCase.run(taskId, width, height, isRtl)
-            ) {
-                is ThumbnailPositionState.MatrixScaling -> thumbnailPositionState.matrix
-                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!!
-        return Snapshot(bitmap, thumbnailData.rotation, task.colorBackground.removeAlpha())
-    }
-
-    @ColorInt private fun Int.removeAlpha(): Int = ColorUtils.setAlphaComponent(this, 0xff)
-
-    private companion object {
-        const val MAX_SCRIM_ALPHA = 0.4f
-        const val TAG = "TaskThumbnailViewModel"
-    }
+    /** Returns a Matrix which can be applied to the snapshot */
+    fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix
 }
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
new file mode 100644
index 0000000..bd47cec
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
@@ -0,0 +1,149 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language goveryning permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.task.viewmodel
+
+import android.annotation.ColorInt
+import android.app.ActivityTaskManager.INVALID_TASK_ID
+import android.graphics.Matrix
+import android.util.Log
+import androidx.core.graphics.ColorUtils
+import com.android.quickstep.recents.data.RecentTasksRepository
+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.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.asStateFlow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.runBlocking
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class TaskThumbnailViewModelImpl(
+    recentsViewData: RecentsViewData,
+    taskViewData: TaskViewData,
+    taskContainerData: TaskContainerData,
+    private val tasksRepository: RecentTasksRepository,
+    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
+
+    override val cornerRadiusProgress =
+        if (taskViewData.isOutlineFormedByThumbnailView) recentsViewData.fullscreenProgress
+        else MutableStateFlow(1f).asStateFlow()
+
+    override val inheritedScale =
+        combine(recentsViewData.scale, taskViewData.scale) { recentsScale, taskScale ->
+            recentsScale * taskScale
+        }
+
+    override val dimProgress: Flow<Float> =
+        combine(taskContainerData.taskMenuOpenProgress, recentsViewData.tintAmount) {
+            taskMenuOpenProgress,
+            tintAmount ->
+            max(taskMenuOpenProgress * MAX_SCRIM_ALPHA, tintAmount)
+        }
+    override val splashAlpha = splashProgress.flatMapLatest { it }
+
+    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 -> LiveTile
+                    isBackgroundOnly(taskVal) ->
+                        BackgroundOnly(taskVal.colorBackground.removeAlpha())
+                    isSnapshotSplashState(taskVal) ->
+                        SnapshotSplash(createSnapshotState(taskVal), taskVal.icon)
+                    else -> Uninitialized
+                }
+            }
+            .distinctUntilChanged()
+
+    override fun bind(taskId: Int) {
+        Log.d(TAG, "bind taskId: $taskId")
+        this.taskId = taskId
+        task.value = tasksRepository.getTaskDataById(taskId)
+        splashProgress.value = splashAlphaUseCase.execute(taskId)
+    }
+
+    override fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix {
+        return runBlocking {
+            when (
+                val thumbnailPositionState =
+                    getThumbnailPositionUseCase.run(taskId, width, height, isRtl)
+            ) {
+                is ThumbnailPositionState.MatrixScaling -> thumbnailPositionState.matrix
+                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!!
+        return Snapshot(bitmap, thumbnailData.rotation, task.colorBackground.removeAlpha())
+    }
+
+    @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/util/AsyncClockEventDelegate.java b/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
index 38ae303..4a84b1b 100644
--- a/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
+++ b/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
@@ -32,23 +32,33 @@
 
 import androidx.annotation.WorkerThread;
 
+import com.android.launcher3.dagger.ApplicationContext;
+import com.android.launcher3.dagger.LauncherAppSingleton;
+import com.android.launcher3.dagger.LauncherBaseAppComponent;
+import com.android.launcher3.util.DaggerSingletonObject;
+import com.android.launcher3.util.DaggerSingletonTracker;
+import com.android.launcher3.util.ExecutorUtil;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.SettingsCache;
 import com.android.launcher3.util.SettingsCache.OnChangeListener;
 import com.android.launcher3.util.SimpleBroadcastReceiver;
+import com.android.quickstep.dagger.QuickstepBaseAppComponent;
 
 import java.util.ArrayList;
 import java.util.List;
 
+import javax.inject.Inject;
+
 /**
  * Extension of {@link ClockEventDelegate} to support async event registration
  */
+@LauncherAppSingleton
 public class AsyncClockEventDelegate extends ClockEventDelegate
         implements OnChangeListener, SafeCloseable {
 
-    public static final MainThreadInitializedObject<AsyncClockEventDelegate> INSTANCE =
-            new MainThreadInitializedObject<>(AsyncClockEventDelegate::new);
+    public static final DaggerSingletonObject<AsyncClockEventDelegate> INSTANCE =
+            new DaggerSingletonObject<>(QuickstepBaseAppComponent::getAsyncClockEventDelegate);
 
     private final Context mContext;
     private final SimpleBroadcastReceiver mReceiver =
@@ -61,10 +71,12 @@
     private boolean mFormatRegistered = false;
     private boolean mDestroyed = false;
 
-    private AsyncClockEventDelegate(Context context) {
+    @Inject
+    AsyncClockEventDelegate(@ApplicationContext Context context, DaggerSingletonTracker tracker) {
         super(context);
         mContext = context;
         mReceiver.register(mContext, ACTION_TIME_CHANGED, ACTION_TIMEZONE_CHANGED);
+        ExecutorUtil.executeSyncOnMainOrFail(() -> tracker.addCloseable(this));
     }
 
     @Override
diff --git a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
index 9335e7e..a5be89a 100644
--- a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
+++ b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
@@ -30,7 +30,6 @@
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
 import android.content.Context;
-import android.content.SharedPreferences;
 import android.graphics.Matrix;
 import android.graphics.Point;
 import android.graphics.PointF;
@@ -45,6 +44,7 @@
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherPrefChangeListener;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.testing.shared.TestProtocol;
 import com.android.launcher3.touch.PagedOrientationHandler;
@@ -66,8 +66,7 @@
  * This class has initial default state assuming the device and foreground app have
  * no ({@link Surface#ROTATION_0} rotation.
  */
-public class RecentsOrientedState implements
-        SharedPreferences.OnSharedPreferenceChangeListener {
+public class RecentsOrientedState implements LauncherPrefChangeListener {
 
     private static final String TAG = "RecentsOrientedState";
     private static final boolean DEBUG = false;
@@ -283,7 +282,7 @@
     }
 
     @Override
-    public void onSharedPreferenceChanged(SharedPreferences sharedPreferences, String s) {
+    public void onPrefChanged(String s) {
         if (LauncherPrefs.ALLOW_ROTATION.getSharedPrefKey().equals(s)) {
             updateHomeRotationSetting();
         }
diff --git a/quickstep/src/com/android/quickstep/util/ScalingWorkspaceRevealAnim.kt b/quickstep/src/com/android/quickstep/util/ScalingWorkspaceRevealAnim.kt
index a94d023..2f0a6df 100644
--- a/quickstep/src/com/android/quickstep/util/ScalingWorkspaceRevealAnim.kt
+++ b/quickstep/src/com/android/quickstep/util/ScalingWorkspaceRevealAnim.kt
@@ -62,7 +62,8 @@
          * Custom interpolator for both the home and wallpaper scaling. Necessary because EMPHASIZED
          * is too aggressive, but EMPHASIZED_DECELERATE is too soft.
          */
-        private val SCALE_INTERPOLATOR =
+        @JvmField
+        val SCALE_INTERPOLATOR =
             PathInterpolator(
                 Path().apply {
                     moveTo(0f, 0f)
diff --git a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
index 00e57c2..bbb8cc8 100644
--- a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -16,7 +16,7 @@
 package com.android.quickstep.views;
 
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
-import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY;
+import static android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY;
 
 import static com.android.launcher3.LauncherState.CLEAR_ALL_BUTTON;
 import static com.android.launcher3.LauncherState.NORMAL;
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index c405080..7554c44 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -129,6 +129,7 @@
 import android.widget.ListView;
 import android.widget.OverScroller;
 import android.widget.Toast;
+import android.window.DesktopModeFlags;
 import android.window.PictureInPictureSurfaceTransaction;
 
 import androidx.annotation.NonNull;
@@ -3935,9 +3936,11 @@
                     if (shouldRemoveTask) {
                         if (dismissedTaskView.isRunningTask()) {
                             finishRecentsAnimation(true /* toRecents */, false /* shouldPip */,
-                                    () -> removeTaskInternal(dismissedTaskViewId));
+                                    () -> removeTaskInternal(dismissedTaskViewId,
+                                            dismissedTaskView instanceof DesktopTaskView));
                         } else {
-                            removeTaskInternal(dismissedTaskViewId);
+                            removeTaskInternal(dismissedTaskViewId,
+                                    dismissedTaskView instanceof DesktopTaskView);
                         }
                         announceForAccessibility(
                                 getResources().getString(R.string.task_view_closed));
@@ -4305,16 +4308,21 @@
         return lastVisibleIndex;
     }
 
-    private void removeTaskInternal(int dismissedTaskViewId) {
+    private void removeTaskInternal(int dismissedTaskViewId, boolean isDesktop) {
         int[] taskIds = getTaskIdsForTaskViewId(dismissedTaskViewId);
-        UI_HELPER_EXECUTOR.getHandler().post(
-                () -> {
-                    for (int taskId : taskIds) {
-                        if (taskId != -1) {
-                            ActivityManagerWrapper.getInstance().removeTask(taskId);
-                        }
+        UI_HELPER_EXECUTOR.getHandler().post(() -> {
+            if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION.isTrue() && isDesktop) {
+                // TODO: b/372005228 - Use the api with desktop id instead.
+                SystemUiProxy.INSTANCE.get(getContext()).removeDesktop(
+                        mContainer.getDisplay().getDisplayId());
+            } else {
+                for (int taskId : taskIds) {
+                    if (taskId != -1) {
+                        ActivityManagerWrapper.getInstance().removeTask(taskId);
                     }
-                });
+                }
+            }
+        });
     }
 
     protected void onDismissAnimationEnds() {
diff --git a/quickstep/src_protolog/com/android/launcher3/util/StateManagerProtoLogProxy.java b/quickstep/src_protolog/com/android/launcher3/util/StateManagerProtoLogProxy.java
new file mode 100644
index 0000000..bc989dc
--- /dev/null
+++ b/quickstep/src_protolog/com/android/launcher3/util/StateManagerProtoLogProxy.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util;
+
+import static com.android.launcher3.Flags.enableStateManagerProtoLog;
+import static com.android.quickstep.util.QuickstepProtoLogGroup.LAUNCHER_STATE_MANAGER;
+
+import androidx.annotation.NonNull;
+
+import com.android.internal.protolog.ProtoLog;
+
+/**
+ * Proxy class used for StateManager ProtoLog support.
+ */
+public class StateManagerProtoLogProxy {
+
+    public static void logGoToState(
+            @NonNull Object fromState, @NonNull Object toState, @NonNull String trace) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER,
+                "StateManager.goToState: fromState: %s, toState: %s, partial trace:\n%s",
+                fromState,
+                toState,
+                trace);
+    }
+
+    public static void logCreateAtomicAnimation(
+            @NonNull Object fromState, @NonNull Object toState, @NonNull String trace) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER, "StateManager.createAtomicAnimation: "
+                        + "fromState: %s, toState: %s, partial trace:\n%s",
+                fromState,
+                toState,
+                trace);
+    }
+
+    public static void logOnStateTransitionStart(@NonNull Object state) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER, "StateManager.onStateTransitionStart: state: %s", state);
+    }
+
+    public static void logOnStateTransitionEnd(@NonNull Object state) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER, "StateManager.onStateTransitionEnd: state: %s", state);
+    }
+
+    public static void logCancelAnimation(boolean animationOngoing, @NonNull String trace) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER,
+                "StateManager.cancelAnimation: animation ongoing: %b, partial trace:\n%s",
+                animationOngoing,
+                trace);
+    }
+}
diff --git a/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java b/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java
index 7b81b9a..bb02a11 100644
--- a/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java
+++ b/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java
@@ -27,7 +27,8 @@
 public enum QuickstepProtoLogGroup implements IProtoLogGroup {
 
     ACTIVE_GESTURE_LOG(true, true, false, "ActiveGestureLog"),
-    RECENTS_WINDOW(true, true, Constants.DEBUG_RECENTS_WINDOW, "RecentsWindow");
+    RECENTS_WINDOW(true, true, Constants.DEBUG_RECENTS_WINDOW, "RecentsWindow"),
+    LAUNCHER_STATE_MANAGER(true, true, Constants.DEBUG_STATE_MANAGER, "LauncherStateManager");
 
     private final boolean mEnabled;
     private volatile boolean mLogToProto;
@@ -97,6 +98,7 @@
     private static final class Constants {
 
         private static final boolean DEBUG_RECENTS_WINDOW = false;
+        private static final boolean DEBUG_STATE_MANAGER = true; // b/279059025, b/325463989
 
         private static final int LOG_START_ID =
                 (int) (UUID.nameUUIDFromBytes(QuickstepProtoLogGroup.class.getName().getBytes())
diff --git a/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt b/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt
index 473d8ef..2f1f0b5 100644
--- a/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt
+++ b/quickstep/testing/com/android/launcher3/taskbar/bubbles/testing/FakeBubbleViewFactory.kt
@@ -73,7 +73,16 @@
                 context.resources.getString(com.android.internal.R.string.config_icon_mask)
             )
         val bubble =
-            BubbleBarBubble(bubbleInfo, bubbleView, badge, icon, dotColor, dotPath, "test app")
+            BubbleBarBubble(
+                bubbleInfo,
+                bubbleView,
+                badge,
+                icon,
+                dotColor,
+                dotPath,
+                "test app",
+                null,
+            )
         bubbleView.setBubble(bubble)
         return bubbleView
     }
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
new file mode 100644
index 0000000..ff5d8bd
--- /dev/null
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.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 cornerRadiusProgress = MutableStateFlow(0f)
+    override val inheritedScale = MutableStateFlow(1f)
+    override val dimProgress = MutableStateFlow(0f)
+    override val splashAlpha = MutableStateFlow(0f)
+    override val uiState = MutableStateFlow<TaskThumbnailUiState>(Uninitialized)
+
+    override fun bind(taskId: Int) {
+        // no-op
+    }
+
+    override fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean) =
+        Matrix.IDENTITY_MATRIX
+}
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
new file mode 100644
index 0000000..75769e9
--- /dev/null
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt
@@ -0,0 +1,86 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.quickstep.task.thumbnail
+
+import android.content.Context
+import android.graphics.Color
+import android.view.LayoutInflater
+import com.android.launcher3.R
+import com.android.quickstep.recents.di.RecentsDependencies
+import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
+import com.google.android.apps.nexuslauncher.imagecomparison.goldenpathmanager.ViewScreenshotGoldenPathManager
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
+import platform.test.screenshot.DeviceEmulationSpec
+import platform.test.screenshot.Displays
+import platform.test.screenshot.ViewScreenshotTestRule
+import platform.test.screenshot.getEmulatedDevicePathConfig
+
+/** Screenshot tests for [TaskThumbnailView]. */
+@RunWith(ParameterizedAndroidJunit4::class)
+class TaskThumbnailViewScreenshotTest(emulationSpec: DeviceEmulationSpec) {
+
+    @get:Rule
+    val screenshotRule =
+        ViewScreenshotTestRule(
+            emulationSpec,
+            ViewScreenshotGoldenPathManager(getEmulatedDevicePathConfig(emulationSpec)),
+        )
+
+    private val taskThumbnailViewModel = FakeTaskThumbnailViewModel()
+
+    @Test
+    fun taskThumbnailView_uninitialized() {
+        screenshotRule.screenshotTest("taskThumbnailView_uninitialized") { activity ->
+            activity.actionBar?.hide()
+            createTaskThumbnailView(activity)
+        }
+    }
+
+    @Test
+    fun taskThumbnailView_backgroundOnly() {
+        screenshotRule.screenshotTest("taskThumbnailView_backgroundOnly") { activity ->
+            activity.actionBar?.hide()
+            taskThumbnailViewModel.uiState.value = TaskThumbnailUiState.BackgroundOnly(Color.YELLOW)
+            createTaskThumbnailView(activity)
+        }
+    }
+
+    private fun createTaskThumbnailView(context: Context): TaskThumbnailView {
+        val di = RecentsDependencies.initialize(context)
+        val taskThumbnailView =
+            LayoutInflater.from(context).inflate(R.layout.task_thumbnail, null, false)
+        val ttvDiScopeId = di.getScope(taskThumbnailView).scopeId
+        di.provide(TaskThumbnailViewData::class.java, ttvDiScopeId) { TaskThumbnailViewData() }
+        di.provide(TaskThumbnailViewModel::class.java, ttvDiScopeId) { taskThumbnailViewModel }
+
+        return taskThumbnailView as TaskThumbnailView
+    }
+
+    companion object {
+        @Parameters(name = "{0}")
+        @JvmStatic
+        fun getTestSpecs() =
+            DeviceEmulationSpec.forDisplays(
+                Displays.Phone,
+                isDarkTheme = false,
+                isLandscape = false,
+            )
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarAutohideSuspendControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarAutohideSuspendControllerTest.kt
index f3fff9f..59900b1 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarAutohideSuspendControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarAutohideSuspendControllerTest.kt
@@ -41,15 +41,14 @@
 @EmulatedDevices(["pixelTablet2023"])
 class TaskbarAutohideSuspendControllerTest {
 
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
-
-    @get:Rule(order = 0) val animatorTestRule = AnimatorTestRule(this)
-    @get:Rule(order = 1)
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val animatorTestRule = AnimatorTestRule(this)
+    @get:Rule(order = 2)
     val systemUiProxyRule = TestRule { base, _ ->
         object : Statement() {
             override fun evaluate() {
                 getInstrumentation().runOnMainSync {
-                    context.applicationContext.putObject(
+                    context.putObject(
                         SystemUiProxy.INSTANCE,
                         object : SystemUiProxy(context) {
                             override fun notifyTaskbarAutohideSuspend(suspend: Boolean) {
@@ -62,8 +61,8 @@
             }
         }
     }
-    @get:Rule(order = 2) val taskbarModeRule = TaskbarModeRule(context)
-    @get:Rule(order = 3) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 3) val taskbarModeRule = TaskbarModeRule(context)
+    @get:Rule(order = 4) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @InjectController lateinit var autohideSuspendController: TaskbarAutohideSuspendController
     @InjectController lateinit var stashController: TaskbarStashController
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarDesktopModeControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarDesktopModeControllerTest.kt
index 72bbfc9..455b6c5 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarDesktopModeControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarDesktopModeControllerTest.kt
@@ -16,7 +16,6 @@
 
 package com.android.launcher3.taskbar
 
-import androidx.test.platform.app.InstrumentationRegistry
 import com.android.launcher3.taskbar.TaskbarBackgroundRenderer.Companion.MAX_ROUNDNESS
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
 import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
@@ -30,12 +29,8 @@
 @LauncherMultivalentJUnit.EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
 class TaskbarDesktopModeControllerTest {
 
-    private val context =
-        TaskbarWindowSandboxContext.create(
-            InstrumentationRegistry.getInstrumentation().targetContext
-        )
-
-    @get:Rule(order = 0) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @TaskbarUnitTestRule.InjectController
     lateinit var taskbarDesktopModeController: TaskbarDesktopModeController
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarEduTooltipControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarEduTooltipControllerTest.kt
index 961d4dc..e575efd 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarEduTooltipControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarEduTooltipControllerTest.kt
@@ -17,7 +17,6 @@
 package com.android.launcher3.taskbar.test
 
 import android.util.Log
-import androidx.test.platform.app.InstrumentationRegistry
 import com.android.launcher3.Utilities
 import com.android.launcher3.taskbar.TOOLTIP_STEP_FEATURES
 import com.android.launcher3.taskbar.TOOLTIP_STEP_NONE
@@ -52,30 +51,21 @@
 @Ignore
 class TaskbarEduTooltipControllerTest {
 
-    private val context =
-        TaskbarWindowSandboxContext.create(
-            InstrumentationRegistry.getInstrumentation().targetContext
-        )
-
-    @get:Rule(order = 0)
-    val tooltipStepPreferenceRule =
-        TaskbarPreferenceRule(
-            context,
-            OnboardingPrefs.TASKBAR_EDU_TOOLTIP_STEP.prefItem,
-        )
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
 
     @get:Rule(order = 1)
+    val tooltipStepPreferenceRule =
+        TaskbarPreferenceRule(context, OnboardingPrefs.TASKBAR_EDU_TOOLTIP_STEP.prefItem)
+
+    @get:Rule(order = 2)
     val searchEduPreferenceRule =
-        TaskbarPreferenceRule(
-            context,
-            OnboardingPrefs.TASKBAR_SEARCH_EDU_SEEN,
-        )
+        TaskbarPreferenceRule(context, OnboardingPrefs.TASKBAR_SEARCH_EDU_SEEN)
 
-    @get:Rule(order = 2) val taskbarPinningPreferenceRule = TaskbarPinningPreferenceRule(context)
+    @get:Rule(order = 3) val taskbarPinningPreferenceRule = TaskbarPinningPreferenceRule(context)
 
-    @get:Rule(order = 3) val taskbarModeRule = TaskbarModeRule(context)
+    @get:Rule(order = 4) val taskbarModeRule = TaskbarModeRule(context)
 
-    @get:Rule(order = 4) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 5) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @InjectController lateinit var taskbarEduTooltipController: TaskbarEduTooltipController
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt
index 3524961..12e84b8 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt
@@ -42,11 +42,10 @@
 @RunWith(LauncherMultivalentJUnit::class)
 @EmulatedDevices(["pixelTablet2023"])
 class TaskbarScrimViewControllerTest {
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
-
-    @get:Rule(order = 0) val taskbarModeRule = TaskbarModeRule(context)
-    @get:Rule(order = 1) val animatorTestRule = AnimatorTestRule(this)
-    @get:Rule(order = 2) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val taskbarModeRule = TaskbarModeRule(context)
+    @get:Rule(order = 2) val animatorTestRule = AnimatorTestRule(this)
+    @get:Rule(order = 3) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @InjectController lateinit var scrimViewController: TaskbarScrimViewController
 
@@ -132,7 +131,7 @@
     @TaskbarMode(PINNED)
     fun testOnClick_scrimShown_performsSystemBack() {
         var backPressed = false
-        context.applicationContext.putObject(
+        context.putObject(
             SystemUiProxy.INSTANCE,
             object : SystemUiProxy(context) {
                 override fun onBackPressed() {
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 d7ce4ed..de73ce7 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
@@ -63,12 +63,11 @@
 @EnableFlags(FLAG_ENABLE_BUBBLE_BAR)
 @EmulatedDevices(["pixelTablet2023"])
 class TaskbarStashControllerTest {
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
-
-    @get:Rule(order = 0) val taskbarModeRule = TaskbarModeRule(context)
-    @get:Rule(order = 1) val taskbarPinningPreferenceRule = TaskbarPinningPreferenceRule(context)
-    @get:Rule(order = 2) val animatorTestRule = AnimatorTestRule(this)
-    @get:Rule(order = 3) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val taskbarModeRule = TaskbarModeRule(context)
+    @get:Rule(order = 2) val taskbarPinningPreferenceRule = TaskbarPinningPreferenceRule(context)
+    @get:Rule(order = 3) val animatorTestRule = AnimatorTestRule(this)
+    @get:Rule(order = 4) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @InjectController lateinit var stashController: TaskbarStashController
     @InjectController lateinit var viewController: TaskbarViewController
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarViewControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarViewControllerTest.kt
index 4aac1dc..b13eafe 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarViewControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarViewControllerTest.kt
@@ -17,7 +17,6 @@
 package com.android.launcher3.taskbar
 
 import android.view.View
-import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.taskbar.TaskbarViewController.DIVIDER_VIEW_POSITION_OFFSET
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
@@ -33,19 +32,23 @@
 @EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
 /**
  * Legend for the comments below:
+ * ```
  * A: All Apps Button
  * H: Hotseat item
  * |: Divider
  * R: Recent item
+ * ```
  *
  * The comments are formatted in two lines:
+ * ```
  * // Items in taskbar, e.g.               A  |  HHHHHH
  * // Index of items relative to Hotseat: -1 -.5 012345
+ * ```
  */
 class TaskbarViewControllerTest {
 
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
-    @get:Rule val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @InjectController lateinit var taskbarViewController: TaskbarViewController
 
@@ -59,7 +62,7 @@
                 /* isAllAppsButton = */ true,
                 /* isTaskbarDividerView = */ false,
                 /* isDividerForRecents = */ false,
-                /* recentTaskIndex = */ -1
+                /* recentTaskIndex = */ -1,
             )
         // [>A<] | [HHHHHH]
         //  -1 -.5  012345
@@ -77,7 +80,7 @@
                 /* isAllAppsButton = */ true,
                 /* isTaskbarDividerView = */ false,
                 /* isDividerForRecents = */ false,
-                /* recentTaskIndex = */ -1
+                /* recentTaskIndex = */ -1,
             )
         // [HHHHHH] | [>A<]
         //  012345 5.5  6
@@ -94,7 +97,7 @@
                 /* isAllAppsButton = */ false,
                 /* isTaskbarDividerView = */ true,
                 /* isDividerForRecents = */ false,
-                /* recentTaskIndex = */ -1
+                /* recentTaskIndex = */ -1,
             )
         // [A] >|< [HHHHHH]
         // -1  -.5  012345
@@ -112,7 +115,7 @@
                 /* isAllAppsButton = */ false,
                 /* isTaskbarDividerView = */ true,
                 /* isDividerForRecents = */ true,
-                /* recentTaskIndex = */ -1
+                /* recentTaskIndex = */ -1,
             )
         // [A] [HHHHHH] >|< [RR]
         // -1   012345  5.5  67
@@ -130,7 +133,7 @@
                 /* isAllAppsButton = */ false,
                 /* isTaskbarDividerView = */ true,
                 /* isDividerForRecents = */ false,
-                /* recentTaskIndex = */ -1
+                /* recentTaskIndex = */ -1,
             )
         // [HHHHHH] >|< [A]
         //  012345  5.5  6
@@ -148,7 +151,7 @@
                 /* isAllAppsButton = */ false,
                 /* isTaskbarDividerView = */ true,
                 /* isDividerForRecents = */ true,
-                /* recentTaskIndex = */ -1
+                /* recentTaskIndex = */ -1,
             )
         // [HHHHHH][A] >|< [RR]
         //  012345  6  6.5  78
@@ -167,7 +170,7 @@
                 /* isAllAppsButton = */ false,
                 /* isTaskbarDividerView = */ false,
                 /* isDividerForRecents = */ false,
-                /* recentTaskIndex = */ recentTaskIndex
+                /* recentTaskIndex = */ recentTaskIndex,
             )
         // [A][HHHHHH] | [>R<R]
         // -1  012345 5.5  6 7
@@ -186,7 +189,7 @@
                 /* isAllAppsButton = */ false,
                 /* isTaskbarDividerView = */ false,
                 /* isDividerForRecents = */ false,
-                /* recentTaskIndex = */ recentTaskIndex
+                /* recentTaskIndex = */ recentTaskIndex,
             )
         // [A][HHHHHH] | [R>R<]
         // -1  012345 5.5 6 7
@@ -205,7 +208,7 @@
                 /* isAllAppsButton = */ false,
                 /* isTaskbarDividerView = */ false,
                 /* isDividerForRecents = */ false,
-                /* recentTaskIndex = */ recentTaskIndex
+                /* recentTaskIndex = */ recentTaskIndex,
             )
         // [HHHHHH][A] | [>R<R]
         //  012345  6 6.5  7 8
@@ -224,7 +227,7 @@
                 /* isAllAppsButton = */ false,
                 /* isTaskbarDividerView = */ false,
                 /* isDividerForRecents = */ false,
-                /* recentTaskIndex = */ recentTaskIndex
+                /* recentTaskIndex = */ recentTaskIndex,
             )
         // [HHHHHH][A] | [R>R<]
         //  012345  6 6.5 7 8
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
index f783e40..60c94a8 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
@@ -44,13 +44,9 @@
 @EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
 class TaskbarAllAppsControllerTest {
 
-    @get:Rule
-    val taskbarUnitTestRule =
-        TaskbarUnitTestRule(
-            this,
-            TaskbarWindowSandboxContext.create(getInstrumentation().targetContext),
-        )
-    @get:Rule val animatorTestRule = AnimatorTestRule(this)
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 2) val animatorTestRule = AnimatorTestRule(this)
 
     @InjectController lateinit var allAppsController: TaskbarAllAppsController
     @InjectController lateinit var overlayController: TaskbarOverlayController
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewControllerTest.kt
index 04f02e9..516220a 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewControllerTest.kt
@@ -47,13 +47,12 @@
 @EmulatedDevices(["pixelFoldable2023"])
 class TaskbarAllAppsViewControllerTest {
 
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
-
-    @get:Rule(order = 0) val taskbarModeRule = TaskbarModeRule(context)
-    @get:Rule(order = 1)
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val taskbarModeRule = TaskbarModeRule(context)
+    @get:Rule(order = 2)
     val allAppsVisitedPreferenceRule =
         TaskbarPreferenceRule(context, ALL_APPS_VISITED_COUNT.prefItem)
-    @get:Rule(order = 2) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 3) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @InjectController lateinit var overlayController: TaskbarOverlayController
     @InjectController lateinit var stashController: TaskbarStashController
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarSwipeControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarSwipeControllerTest.kt
index 3b6952d..2e471b8 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarSwipeControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleBarSwipeControllerTest.kt
@@ -37,6 +37,7 @@
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.atLeastOnce
 import org.mockito.kotlin.never
+import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
 
@@ -45,15 +46,11 @@
 
     companion object {
         const val UNSTASH_THRESHOLD = 100
-        const val EXPAND_THRESHOLD = 200
         const val MAX_OVERSCROLL = 300
-        const val STASH_THRESHOLD = 50
 
         const val UP_BELOW_UNSTASH = -UNSTASH_THRESHOLD + 10f
         const val UP_ABOVE_UNSTASH = -UNSTASH_THRESHOLD - 10f
-        const val UP_ABOVE_EXPAND = -EXPAND_THRESHOLD - 10f
-        const val DOWN_UNDER_STASH = STASH_THRESHOLD - 10f
-        const val DOWN_OVER_STASH = STASH_THRESHOLD + 10f
+        const val DOWN = UNSTASH_THRESHOLD + 10f
     }
 
     private val context = ApplicationProvider.getApplicationContext<Context>()
@@ -80,14 +77,8 @@
                 override val unstashThreshold: Int
                     get() = UNSTASH_THRESHOLD
 
-                override val expandThreshold: Int
-                    get() = EXPAND_THRESHOLD
-
                 override val maxOverscroll: Int
                     get() = MAX_OVERSCROLL
-
-                override val stashThreshold: Int
-                    get() = STASH_THRESHOLD
             }
         bubbleBarSwipeController = BubbleBarSwipeController(context, dimensionProvider)
 
@@ -135,35 +126,11 @@
     }
 
     @Test
-    fun swipeUp_stashedBar_aboveExpandThreshold_viewsHaveDampedTranslation() {
-        setUpStashedBar()
-        testViewsHaveDampedTranslationOnSwipe(UP_ABOVE_EXPAND)
-    }
-
-    @Test
     fun swipeUp_collapsedBar_aboveUnstashThreshold_viewsHaveDampedTranslation() {
         setUpCollapsedBar()
         testViewsHaveDampedTranslationOnSwipe(UP_ABOVE_UNSTASH)
     }
 
-    @Test
-    fun swipeUp_collapsedBar_aboveExpandThreshold_viewsHaveDampedTranslation() {
-        setUpCollapsedBar()
-        testViewsHaveDampedTranslationOnSwipe(UP_ABOVE_EXPAND)
-    }
-
-    @Test
-    fun swipeDown_collapsedBar_belowStashThreshold_viewsHaveDampedTranslation() {
-        setUpCollapsedBar()
-        testViewsHaveDampedTranslationOnSwipe(DOWN_UNDER_STASH)
-    }
-
-    @Test
-    fun swipeDown_collapsedBar_overStashThreshold_viewsHaveDampedTranslation() {
-        setUpCollapsedBar()
-        testViewsHaveDampedTranslationOnSwipe(DOWN_OVER_STASH)
-    }
-
     // endregion
 
     // region Test that translation on views is reset on finish
@@ -203,29 +170,11 @@
     }
 
     @Test
-    fun swipeUp_stashedBar_aboveExpandThreshold_animateTranslationToZeroOnFinish() {
-        setUpStashedBar()
-        testViewsTranslationResetOnFinish(UP_ABOVE_EXPAND)
-    }
-
-    @Test
     fun swipeUp_collapsedBar_aboveUnstashThreshold_animateTranslationToZeroOnFinish() {
         setUpCollapsedBar()
         testViewsTranslationResetOnFinish(UP_ABOVE_UNSTASH)
     }
 
-    @Test
-    fun swipeUp_collapsedBar_aboveExpandThreshold_animateTranslationToZeroOnFinish() {
-        setUpCollapsedBar()
-        testViewsTranslationResetOnFinish(UP_ABOVE_EXPAND)
-    }
-
-    @Test
-    fun swipeDown_collapsedBar_aboveStashThreshold_animateTranslationToZeroOnFinish() {
-        setUpCollapsedBar()
-        testViewsTranslationResetOnFinish(DOWN_OVER_STASH)
-    }
-
     // endregion
 
     // region Test swipe interactions on stashed bar
@@ -251,7 +200,7 @@
     }
 
     @Test
-    fun swipeUp_stashedBar_aboveUnstashThreshold_unstashBubbleBar() {
+    fun swipeUp_stashedBar_overUnstashThreshold_unstashBubbleBar() {
         setUpStashedBar()
         getInstrumentation().runOnMainSync {
             bubbleBarSwipeController.start()
@@ -271,50 +220,45 @@
     }
 
     @Test
-    fun swipeUp_stashedBar_overUnstashThresholdMultipleTimes_unstashBubbleBarOnce() {
+    fun swipeUp_stashedBar_overUnstashThresholdMultipleTimes_unstashesMultipleTimes() {
         setUpStashedBar()
         getInstrumentation().runOnMainSync {
             bubbleBarSwipeController.start()
             bubbleBarSwipeController.swipeTo(UP_ABOVE_UNSTASH)
             bubbleBarSwipeController.swipeTo(UP_BELOW_UNSTASH)
-            bubbleBarSwipeController.swipeTo(UP_ABOVE_UNSTASH)
         }
         verify(bubbleStashController).showBubbleBar(expandBubbles = false)
+        verify(bubbleStashController).stashBubbleBar()
+
+        getInstrumentation().runOnMainSync { bubbleBarSwipeController.swipeTo(UP_ABOVE_UNSTASH) }
+        verify(bubbleStashController, times(2)).showBubbleBar(expandBubbles = false)
     }
 
     @Test
-    fun swipeUp_stashedBar_overExpandThreshold_doesNotExpandBeforeFinish() {
+    fun swipeUp_stashedBar_releaseOverUnstashThreshold_expandsBar() {
         setUpStashedBar()
         getInstrumentation().runOnMainSync {
             bubbleBarSwipeController.start()
-            bubbleBarSwipeController.swipeTo(UP_ABOVE_EXPAND)
+            bubbleBarSwipeController.swipeTo(UP_ABOVE_UNSTASH)
         }
-        verify(bubbleStashController).showBubbleBar(expandBubbles = false)
+        verify(bubbleStashController, never()).showBubbleBar(expandBubbles = true)
         getInstrumentation().runOnMainSync { bubbleBarSwipeController.finish() }
         verify(bubbleStashController).showBubbleBar(expandBubbles = true)
     }
 
     @Test
-    fun swipeUp_stashedBar_overExpandThreshold_isSwipeGestureTrue() {
+    fun swipeUp_stashedBar_overUnstashReleaseBelowUnstash_doesNotExpandBar() {
         setUpStashedBar()
         getInstrumentation().runOnMainSync {
             bubbleBarSwipeController.start()
-            bubbleBarSwipeController.swipeTo(UP_ABOVE_EXPAND)
-        }
-        assertThat(bubbleBarSwipeController.isSwipeGesture()).isTrue()
-    }
-
-    @Test
-    fun swipeUp_stashedBar_overExpandThresholdAndBackDown_doesNotExpandAfterFinish() {
-        setUpStashedBar()
-        getInstrumentation().runOnMainSync {
-            bubbleBarSwipeController.start()
-            bubbleBarSwipeController.swipeTo(UP_ABOVE_EXPAND)
             bubbleBarSwipeController.swipeTo(UP_ABOVE_UNSTASH)
         }
         verify(bubbleStashController).showBubbleBar(expandBubbles = false)
-        getInstrumentation().runOnMainSync { bubbleBarSwipeController.finish() }
-        verify(bubbleStashController).showBubbleBar(expandBubbles = false)
+        getInstrumentation().runOnMainSync {
+            bubbleBarSwipeController.swipeTo(UP_BELOW_UNSTASH)
+            bubbleBarSwipeController.finish()
+        }
+        verify(bubbleStashController, never()).showBubbleBar(expandBubbles = true)
     }
 
     @Test
@@ -322,7 +266,7 @@
         setUpStashedBar()
         getInstrumentation().runOnMainSync {
             bubbleBarSwipeController.start()
-            bubbleBarSwipeController.swipeTo(DOWN_OVER_STASH)
+            bubbleBarSwipeController.swipeTo(DOWN)
         }
         verify(bubbleStashedHandleViewController, never()).setTranslationYForSwipe(any())
         verify(bubbleBarViewController, never()).setTranslationYForSwipe(any())
@@ -338,8 +282,8 @@
         setUpExpandedBar()
         getInstrumentation().runOnMainSync {
             bubbleBarSwipeController.start()
-            bubbleBarSwipeController.swipeTo(UP_ABOVE_EXPAND)
-            bubbleBarSwipeController.swipeTo(DOWN_OVER_STASH)
+            bubbleBarSwipeController.swipeTo(UP_ABOVE_UNSTASH)
+            bubbleBarSwipeController.swipeTo(DOWN)
             bubbleBarSwipeController.finish()
         }
         verify(bubbleStashedHandleViewController, never()).setTranslationYForSwipe(any())
@@ -352,48 +296,71 @@
     // region Test swipe interactions on collapsed bar
 
     @Test
-    fun swipeDown_collapsedBar_underStashThreshold_doesNotHideBar() {
+    fun swipeUp_collapsedBar_doesNotShowBarDuringDrag() {
         setUpCollapsedBar()
         getInstrumentation().runOnMainSync {
             bubbleBarSwipeController.start()
-            bubbleBarSwipeController.swipeTo(DOWN_UNDER_STASH)
-            bubbleBarSwipeController.finish()
+            bubbleBarSwipeController.swipeTo(UP_BELOW_UNSTASH)
+            bubbleBarSwipeController.swipeTo(UP_ABOVE_UNSTASH)
         }
-        verify(bubbleStashController, never()).stashBubbleBar()
+        verify(bubbleStashController, never()).showBubbleBar(any())
     }
 
     @Test
-    fun swipeDown_collapsedBar_overStashThreshold_doesNotHideBarBeforeFinish() {
+    fun swipeUp_collapsedBar_belowUnstashThreshold_isSwipeGestureFalse() {
         setUpCollapsedBar()
         getInstrumentation().runOnMainSync {
             bubbleBarSwipeController.start()
-            bubbleBarSwipeController.swipeTo(DOWN_OVER_STASH)
-        }
-        verify(bubbleStashController, never()).stashBubbleBar()
-        getInstrumentation().runOnMainSync { bubbleBarSwipeController.finish() }
-        verify(bubbleStashController).stashBubbleBar()
-    }
-
-    @Test
-    fun swipeDown_collapsedBar_underStashThreshold_isSwipeGestureFalse() {
-        setUpCollapsedBar()
-        getInstrumentation().runOnMainSync {
-            bubbleBarSwipeController.start()
-            bubbleBarSwipeController.swipeTo(DOWN_UNDER_STASH)
+            bubbleBarSwipeController.swipeTo(UP_BELOW_UNSTASH)
         }
         assertThat(bubbleBarSwipeController.isSwipeGesture()).isFalse()
     }
 
     @Test
-    fun swipeDown_collapsedBar_overStashThreshold_isSwipeGestureTrue() {
+    fun swipeUp_collapsedBar_overUnstashThreshold_isSwipeGestureTrue() {
         setUpCollapsedBar()
         getInstrumentation().runOnMainSync {
             bubbleBarSwipeController.start()
-            bubbleBarSwipeController.swipeTo(DOWN_OVER_STASH)
+            bubbleBarSwipeController.swipeTo(UP_ABOVE_UNSTASH)
         }
         assertThat(bubbleBarSwipeController.isSwipeGesture()).isTrue()
     }
 
+    @Test
+    fun swipeUp_collapsedBar_finishOverUnstashThreshold_expandsBar() {
+        setUpCollapsedBar()
+        getInstrumentation().runOnMainSync {
+            bubbleBarSwipeController.start()
+            bubbleBarSwipeController.swipeTo(UP_ABOVE_UNSTASH)
+            bubbleBarSwipeController.finish()
+        }
+        verify(bubbleStashController).showBubbleBar(expandBubbles = true)
+    }
+
+    @Test
+    fun swipeUp_collapsedBar_finishBelowUnstashThreshold_doesNotExpandBar() {
+        setUpCollapsedBar()
+        getInstrumentation().runOnMainSync {
+            bubbleBarSwipeController.start()
+            bubbleBarSwipeController.swipeTo(UP_BELOW_UNSTASH)
+            bubbleBarSwipeController.finish()
+        }
+        verify(bubbleStashController, never()).showBubbleBar(any())
+    }
+
+    @Test
+    fun swipeDown_collapsedBar_swipeIgnored() {
+        setUpCollapsedBar()
+        getInstrumentation().runOnMainSync {
+            bubbleBarSwipeController.start()
+            bubbleBarSwipeController.swipeTo(DOWN)
+        }
+        verify(bubbleStashedHandleViewController, never()).setTranslationYForSwipe(any())
+        verify(bubbleBarViewController, never()).setTranslationYForSwipe(any())
+        verify(bubbleStashController, never()).showBubbleBar(any())
+        verify(bubbleStashController, never()).stashBubbleBar()
+    }
+
     // endregion
 
     private fun setUpStashedBar() {
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt
index 2caff01..4ae8877 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt
@@ -80,7 +80,16 @@
                 )
             bubbleView = inflater.inflate(R.layout.bubblebar_item_view, null, false) as BubbleView
             bubble =
-                BubbleBarBubble(bubbleInfo, bubbleView, bitmap, bitmap, Color.WHITE, Path(), "")
+                BubbleBarBubble(
+                    bubbleInfo,
+                    bubbleView,
+                    bitmap,
+                    bitmap,
+                    Color.WHITE,
+                    Path(),
+                    "",
+                    null,
+                )
             bubbleView.setBubble(bubble)
         }
         InstrumentationRegistry.getInstrumentation().waitForIdleSync()
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
index 3913165..7eee4de 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
@@ -886,7 +886,16 @@
             bubbleView =
                 inflater.inflate(R.layout.bubblebar_item_view, bubbleBarView, false) as BubbleView
             bubble =
-                BubbleBarBubble(bubbleInfo, bubbleView, bitmap, bitmap, Color.WHITE, Path(), "")
+                BubbleBarBubble(
+                    bubbleInfo,
+                    bubbleView,
+                    bitmap,
+                    bitmap,
+                    Color.WHITE,
+                    Path(),
+                    "",
+                    null,
+                )
             bubbleView.setBubble(bubble)
             bubbleBarView.addView(bubbleView)
         }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
index 4fa821d..1113129 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
@@ -18,7 +18,6 @@
 
 import android.app.ActivityManager.RunningTaskInfo
 import android.view.MotionEvent
-import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingView.TYPE_OPTIONS_POPUP
 import com.android.launcher3.AbstractFloatingView.TYPE_TASKBAR_ALL_APPS
@@ -42,12 +41,8 @@
 @EmulatedDevices(["pixelFoldable2023"])
 class TaskbarOverlayControllerTest {
 
-    @get:Rule
-    val taskbarUnitTestRule =
-        TaskbarUnitTestRule(
-            this,
-            TaskbarWindowSandboxContext.create(getInstrumentation().targetContext),
-        )
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
     @InjectController lateinit var overlayController: TaskbarOverlayController
 
     private val taskbarContext: TaskbarActivityContext
@@ -223,9 +218,8 @@
     }
 
     private class TestOverlayView
-    private constructor(
-        private val overlayContext: TaskbarOverlayContext,
-    ) : AbstractFloatingView(overlayContext, null) {
+    private constructor(private val overlayContext: TaskbarOverlayContext) :
+        AbstractFloatingView(overlayContext, null) {
 
         var type = TYPE_OPTIONS_POPUP
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
index c48947e..74b154a 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
@@ -61,7 +61,7 @@
                 val mode = taskbarMode.mode
 
                 getInstrumentation().runOnMainSync {
-                    context.applicationContext.putObject(
+                    context.putObject(
                         DisplayController.INSTANCE,
                         object : DisplayController(context) {
                             override fun getInfo(): Info {
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt
index f7e4576..0dd1324 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt
@@ -16,7 +16,6 @@
 
 package com.android.launcher3.taskbar.rules
 
-import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.PINNED
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.THREE_BUTTONS
@@ -35,9 +34,8 @@
 @EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
 class TaskbarModeRuleTest {
 
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
-
-    @get:Rule val taskbarModeRule = TaskbarModeRule(context)
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val taskbarModeRule = TaskbarModeRule(context)
 
     @Test
     @TaskbarMode(TRANSIENT)
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt
index a515405..977e7a5 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt
@@ -16,13 +16,13 @@
 
 package com.android.launcher3.taskbar.rules
 
-import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.android.launcher3.util.window.WindowManagerProxy
 import com.google.android.apps.nexuslauncher.deviceemulator.TestWindowManagerProxy
 import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.Description
 import org.junit.runner.RunWith
@@ -31,7 +31,7 @@
 @RunWith(LauncherMultivalentJUnit::class)
 @EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
 class TaskbarPinningPreferenceRuleTest {
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
+    @get:Rule val context = TaskbarWindowSandboxContext.create()
 
     private val preferenceRule = TaskbarPinningPreferenceRule(context)
 
@@ -55,7 +55,7 @@
 
     @Test
     fun testEnableDesktopPinning_verifyDisplayController() {
-        context.applicationContext.putObject(
+        context.putObject(
             WindowManagerProxy.INSTANCE,
             TestWindowManagerProxy(context).apply { isInDesktopMode = true },
         )
@@ -69,7 +69,7 @@
 
     @Test
     fun testDisableDesktopPinning_verifyDisplayController() {
-        context.applicationContext.putObject(
+        context.putObject(
             WindowManagerProxy.INSTANCE,
             TestWindowManagerProxy(context).apply { isInDesktopMode = true },
         )
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRule.kt
index a76a77d..e42ca9e 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRule.kt
@@ -29,11 +29,12 @@
  * The original preference value is restored on teardown.
  */
 class TaskbarPreferenceRule<T : Any>(
-    context: TaskbarWindowSandboxContext,
-    private val constantItem: ConstantItem<T>
+    private val context: TaskbarWindowSandboxContext,
+    private val constantItem: ConstantItem<T>,
 ) : TestRule {
 
-    private val prefs = LauncherPrefs.get(context)
+    private val prefs: LauncherPrefs
+        get() = LauncherPrefs.get(context)
 
     var value: T
         get() = prefs.get(constantItem)
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRuleTest.kt
index 46817d2..b7e6fa3 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRuleTest.kt
@@ -16,11 +16,11 @@
 
 package com.android.launcher3.taskbar.rules
 
-import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.LauncherPrefs.Companion.TASKBAR_PINNING
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.Description
 import org.junit.runner.RunWith
@@ -30,7 +30,7 @@
 @EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
 class TaskbarPreferenceRuleTest {
 
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
+    @get:Rule val context = TaskbarWindowSandboxContext.create()
     private val preferenceRule = TaskbarPreferenceRule(context, TASKBAR_PINNING)
 
     @Test
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
index 52ca78d..7daa142 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
@@ -19,7 +19,6 @@
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
-import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.taskbar.TaskbarKeyguardController
 import com.android.launcher3.taskbar.TaskbarManager
@@ -43,9 +42,8 @@
 @EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
 class TaskbarUnitTestRuleTest {
 
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
-
-    @get:Rule(order = 0) val setFlagsRule = SetFlagsRule()
+    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 1) val setFlagsRule = SetFlagsRule()
 
     @Test
     fun testSetup_taskbarInitialized() {
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
index ee21df8..7728504 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
@@ -18,44 +18,59 @@
 
 import android.content.Context
 import android.content.ContextWrapper
-import android.os.Bundle
-import android.view.Display
-import com.android.launcher3.util.MainThreadInitializedObject
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
+import android.hardware.display.DisplayManager
+import android.hardware.display.VirtualDisplay
+import android.view.Display.DEFAULT_DISPLAY
+import androidx.test.core.app.ApplicationProvider
+import com.android.launcher3.util.MainThreadInitializedObject.ObjectSandbox
+import com.android.launcher3.util.SandboxApplication
+import org.junit.rules.ExternalResource
+import org.junit.rules.RuleChain
+import org.junit.rules.TestRule
 
 /**
- * Sandbox wrapper where [createWindowContext] provides contexts that are still sandboxed within
- * [application].
+ * [SandboxApplication] for running Taskbar tests.
  *
- * Taskbar can create window contexts, which need to operate under the same sandbox application, but
- * [Context.getApplicationContext] by default returns the actual application. For this reason,
- * [SandboxContext] overrides [getApplicationContext] to return itself, which prevents leaving the
- * sandbox. [SandboxContext] and the real application have different sets of
- * [MainThreadInitializedObject] instances, so overriding the application prevents the latter set
- * from leaking into the sandbox. Similarly, this implementation overrides [getApplicationContext]
- * to return the original sandboxed [application], and it wraps created windowed contexts to
- * propagate this [application].
+ * Tests need to run on a [VirtualDisplay] to avoid conflicting with Launcher's Taskbar on the
+ * [DEFAULT_DISPLAY] (i.e. test is executing on a device).
  */
 class TaskbarWindowSandboxContext
-private constructor(private val application: SandboxContext, base: Context) : ContextWrapper(base) {
-
-    override fun createWindowContext(type: Int, options: Bundle?): Context {
-        return TaskbarWindowSandboxContext(application, super.createWindowContext(type, options))
-    }
-
-    override fun createWindowContext(display: Display, type: Int, options: Bundle?): Context {
-        return TaskbarWindowSandboxContext(
-            application,
-            super.createWindowContext(display, type, options),
-        )
-    }
-
-    override fun getApplicationContext(): SandboxContext = application
+private constructor(base: SandboxApplication, val virtualDisplay: VirtualDisplay) :
+    ContextWrapper(base),
+    ObjectSandbox by base,
+    TestRule by RuleChain.outerRule(virtualDisplayRule(virtualDisplay)).around(base) {
 
     companion object {
-        /** Creates a [TaskbarWindowSandboxContext] to sandbox [base] for Taskbar tests. */
-        fun create(base: Context): TaskbarWindowSandboxContext {
-            return SandboxContext(base).let { TaskbarWindowSandboxContext(it, it) }
+        private const val VIRTUAL_DISPLAY_NAME = "TaskbarSandboxDisplay"
+
+        /** Creates a [SandboxApplication] for Taskbar tests. */
+        fun create(): TaskbarWindowSandboxContext {
+            val base = ApplicationProvider.getApplicationContext<Context>()
+            val displayManager = checkNotNull(base.getSystemService(DisplayManager::class.java))
+
+            // Create virtual display to avoid clashing with Taskbar on default display.
+            val virtualDisplay =
+                base.resources.displayMetrics.let {
+                    displayManager.createVirtualDisplay(
+                        VIRTUAL_DISPLAY_NAME,
+                        it.widthPixels,
+                        it.heightPixels,
+                        it.densityDpi,
+                        /* surface= */ null,
+                        /* flags= */ 0,
+                    )
+                }
+
+            return TaskbarWindowSandboxContext(
+                SandboxApplication(base.createDisplayContext(virtualDisplay.display)),
+                virtualDisplay,
+            )
         }
     }
 }
+
+private fun virtualDisplayRule(virtualDisplay: VirtualDisplay): TestRule {
+    return object : ExternalResource() {
+        override fun after() = virtualDisplay.release()
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
index 4834d48..69095e7 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
@@ -16,30 +16,32 @@
 
 package com.android.launcher3.taskbar.rules
 
-import android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY
-import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.util.LauncherMultivalentJUnit
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
+import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.google.common.truth.Truth.assertThat
 import org.junit.Test
+import org.junit.runner.Description
 import org.junit.runner.RunWith
+import org.junit.runners.model.Statement
 
 @RunWith(LauncherMultivalentJUnit::class)
-@LauncherMultivalentJUnit.EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
+@EmulatedDevices(["pixelFoldable2023"])
 class TaskbarWindowSandboxContextTest {
 
-    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
-
     @Test
-    fun testCreateWindowContext_applicationContextSandboxed() {
-        val windowContext = context.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
-        assertThat(windowContext.applicationContext).isInstanceOf(SandboxContext::class.java)
-    }
+    fun testVirtualDisplay_releasedOnTeardown() {
+        val context = TaskbarWindowSandboxContext.create()
+        assertThat(context.virtualDisplay.token).isNotNull()
 
-    @Test
-    fun testCreateWindowContext_nested_applicationContextSandboxed() {
-        val windowContext = context.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
-        val nestedContext = windowContext.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
-        assertThat(nestedContext.applicationContext).isInstanceOf(SandboxContext::class.java)
+        context
+            .apply(
+                object : Statement() {
+                    override fun evaluate() = Unit
+                },
+                Description.createSuiteDescription(TaskbarWindowSandboxContextTest::class.java),
+            )
+            .evaluate()
+
+        assertThat(context.virtualDisplay.token).isNull()
     }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
index a87465f..d55f2e3 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
@@ -22,7 +22,6 @@
 import android.graphics.drawable.Drawable
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.launcher3.util.TestDispatcherProvider
-import com.android.quickstep.task.thumbnail.TaskThumbnailViewModelTest
 import com.android.quickstep.util.DesktopTask
 import com.android.quickstep.util.GroupTask
 import com.android.systemui.shared.recents.model.Task
@@ -286,9 +285,14 @@
 
     private fun createThumbnailData(): ThumbnailData {
         val bitmap = mock<Bitmap>()
-        whenever(bitmap.width).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_WIDTH)
-        whenever(bitmap.height).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_HEIGHT)
+        whenever(bitmap.width).thenReturn(THUMBNAIL_WIDTH)
+        whenever(bitmap.height).thenReturn(THUMBNAIL_HEIGHT)
 
         return ThumbnailData(thumbnail = bitmap)
     }
+
+    companion object {
+        const val THUMBNAIL_WIDTH = 100
+        const val THUMBNAIL_HEIGHT = 200
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
index 33d96a8..829987c 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
@@ -22,7 +22,6 @@
 import android.graphics.Color
 import android.view.Surface
 import com.android.quickstep.recents.data.FakeTasksRepository
-import com.android.quickstep.task.thumbnail.TaskThumbnailViewModelTest
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
@@ -126,9 +125,14 @@
 
     private fun createThumbnailData(rotation: Int = Surface.ROTATION_0): ThumbnailData {
         val bitmap = mock<Bitmap>()
-        whenever(bitmap.width).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_WIDTH)
-        whenever(bitmap.height).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_HEIGHT)
+        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
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
similarity index 97%
rename from quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
rename to quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
index fcf4e56..c88a3fc 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
@@ -36,7 +36,7 @@
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.SnapshotSplash
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
 import com.android.quickstep.task.viewmodel.TaskContainerData
-import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
+import com.android.quickstep.task.viewmodel.TaskThumbnailViewModelImpl
 import com.android.quickstep.task.viewmodel.TaskViewData
 import com.android.quickstep.views.TaskViewType
 import com.android.systemui.shared.recents.model.Task
@@ -51,7 +51,7 @@
 
 /** Test for [TaskThumbnailView] */
 @RunWith(AndroidJUnit4::class)
-class TaskThumbnailViewModelTest {
+class TaskThumbnailViewModelImplTest {
     private var taskViewType = TaskViewType.SINGLE
     private val recentsViewData = RecentsViewData()
     private val taskViewData by lazy { TaskViewData(taskViewType) }
@@ -60,7 +60,7 @@
     private val mGetThumbnailPositionUseCase = mock<GetThumbnailPositionUseCase>()
     private val splashAlphaUseCase: SplashAlphaUseCase = mock()
     private val systemUnderTest by lazy {
-        TaskThumbnailViewModel(
+        TaskThumbnailViewModelImpl(
             recentsViewData,
             taskViewData,
             taskContainerData,
@@ -109,7 +109,7 @@
                         bitmap = expectedThumbnailData.thumbnail!!,
                         thumbnailRotation = Surface.ROTATION_0,
                     ),
-                    expectedIconData.icon
+                    expectedIconData.icon,
                 )
             )
     }
@@ -204,7 +204,7 @@
                         bitmap = expectedThumbnailData.thumbnail!!,
                         thumbnailRotation = Surface.ROTATION_270,
                     ),
-                    expectedIconData.icon
+                    expectedIconData.icon,
                 )
             )
     }
@@ -230,7 +230,7 @@
                         bitmap = expectedThumbnailData.thumbnail!!,
                         thumbnailRotation = Surface.ROTATION_0,
                     ),
-                    expectedIconData.icon
+                    expectedIconData.icon,
                 )
             )
     }
diff --git a/res/drawable/ic_private_profile_divider_badge.xml b/res/drawable/ic_private_profile_divider_badge.xml
new file mode 100644
index 0000000..07c740d
--- /dev/null
+++ b/res/drawable/ic_private_profile_divider_badge.xml
@@ -0,0 +1,26 @@
+<!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="20dp"
+    android:height="20dp"
+    android:viewportWidth="20"
+    android:viewportHeight="20">
+    <group>
+      <path
+          android:pathData="M5,9L15,9A1,1 0,0 1,16 10L16,10A1,1 0,0 1,15 11L5,11A1,1 0,0 1,4 10L4,10A1,1 0,0 1,5 9z"
+          android:fillColor="?attr/materialColorOnSurface"/>
+    </group>
+</vector>
diff --git a/res/drawable/ic_private_profile_letter_list_fast_scroller_badge.xml b/res/drawable/ic_private_profile_letter_list_fast_scroller_badge.xml
new file mode 100644
index 0000000..8d12598
--- /dev/null
+++ b/res/drawable/ic_private_profile_letter_list_fast_scroller_badge.xml
@@ -0,0 +1,28 @@
+<!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="12dp"
+    android:height="15dp"
+    android:viewportWidth="12"
+    android:viewportHeight="15">
+    <path
+        android:pathData="M5.952,0.911L0.645,2.902V6.942C0.645,10.292 2.907,13.417 5.952,14.18C8.997,13.417 11.26,10.292 11.26,6.942V2.902L5.952,0.911ZM7.943,9.536V10.863H6.616V11.526H5.289V8.103C4.333,7.818 3.63,6.942 3.63,5.887C3.63,4.607 4.672,3.565 5.952,3.565C7.233,3.565 8.274,4.607 8.274,5.887C8.274,6.935 7.571,7.818 6.616,8.103V9.536H7.943Z"
+        android:fillColor="#3C4043"
+        android:fillType="evenOdd"/>
+    <path
+        android:pathData="M5.952,6.882C6.502,6.882 6.947,6.436 6.947,5.887C6.947,5.337 6.502,4.892 5.952,4.892C5.403,4.892 4.957,5.337 4.957,5.887C4.957,6.436 5.403,6.882 5.952,6.882Z"
+        android:fillColor="#3C4043"/>
+</vector>
diff --git a/res/layout/user_folder_icon_normalized.xml b/res/layout/user_folder_icon_normalized.xml
index 43a8aac..002e7b7 100644
--- a/res/layout/user_folder_icon_normalized.xml
+++ b/res/layout/user_folder_icon_normalized.xml
@@ -40,12 +40,11 @@
         <com.android.launcher3.folder.FolderNameEditText
             android:id="@+id/folder_name"
             android:layout_width="0dp"
-            android:layout_height="wrap_content"
-            android:layout_gravity="center_vertical"
+            android:layout_height="match_parent"
             style="@style/TextHeadline"
             android:layout_weight="1"
             android:background="@android:color/transparent"
-            android:gravity="center_horizontal"
+            android:gravity="center"
             android:hint="@string/folder_hint_text"
             android:imeOptions="flagNoExtractUi"
             android:importantForAutofill="no"
diff --git a/res/values-gl/strings.xml b/res/values-gl/strings.xml
index e6d6442..02389c4 100644
--- a/res/values-gl/strings.xml
+++ b/res/values-gl/strings.xml
@@ -120,7 +120,7 @@
     <string name="app_pair_name_format" msgid="8134106404716224054">"Emparellamento de aplicacións: <xliff:g id="APP1">%1$s</xliff:g> e <xliff:g id="APP2">%2$s</xliff:g>"</string>
     <string name="styles_wallpaper_button_text" msgid="8216961355289236794">"Estilo e fondo de pantalla"</string>
     <string name="edit_home_screen" msgid="8947858375782098427">"Editar pantalla de inicio"</string>
-    <string name="settings_button_text" msgid="8873672322605444408">"Axustes de Inicio"</string>
+    <string name="settings_button_text" msgid="8873672322605444408">"Configuración da pantalla de inicio"</string>
     <string name="msg_disabled_by_admin" msgid="6898038085516271325">"Función desactivada polo administrador"</string>
     <string name="allow_rotation_title" msgid="7222049633713050106">"Permitir xirar a pantalla de inicio"</string>
     <string name="allow_rotation_desc" msgid="8662546029078692509">"Ao xirar o teléfono"</string>
diff --git a/res/values-kn/strings.xml b/res/values-kn/strings.xml
index 91530b5..0c46eb4 100644
--- a/res/values-kn/strings.xml
+++ b/res/values-kn/strings.xml
@@ -120,7 +120,7 @@
     <string name="app_pair_name_format" msgid="8134106404716224054">"ಆ್ಯಪ್ ಜೋಡಿ: <xliff:g id="APP1">%1$s</xliff:g> ಮತ್ತು <xliff:g id="APP2">%2$s</xliff:g>"</string>
     <string name="styles_wallpaper_button_text" msgid="8216961355289236794">"ವಾಲ್‌ಪೇಪರ್ ಮತ್ತು ಶೈಲಿ"</string>
     <string name="edit_home_screen" msgid="8947858375782098427">"ಹೋಮ್ ಸ್ಕ್ರೀನ್ ಅನ್ನು ಎಡಿಟ್ ಮಾಡಿ"</string>
-    <string name="settings_button_text" msgid="8873672322605444408">"ಮುಖಪುಟ ಸೆಟ್ಟಿಂಗ್‌ಗಳು"</string>
+    <string name="settings_button_text" msgid="8873672322605444408">"ಹೋಮ್ ಸೆಟ್ಟಿಂಗ್‌ಗಳು"</string>
     <string name="msg_disabled_by_admin" msgid="6898038085516271325">"ನಿಮ್ಮ ನಿರ್ವಾಹಕರು ನಿಷ್ಕ್ರಿಯಗೊಳಿಸಿದ್ದಾರೆ"</string>
     <string name="allow_rotation_title" msgid="7222049633713050106">"ಹೋಮ್ ಸ್ಕ್ರೀನ್ ತಿರುಗುವಿಕೆಯನ್ನು ಅನುಮತಿಸಿ"</string>
     <string name="allow_rotation_desc" msgid="8662546029078692509">"ಫೋನ್‌ ತಿರುಗಿಸಿದಾಗ"</string>
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index b0ec9b0..a8840fe 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -412,6 +412,7 @@
 
     private final List<BackPressHandler> mBackPressedHandlers = new ArrayList<>();
     private boolean mIsColdStartupAfterReboot;
+    private boolean mForceConfigUpdate;
 
     private boolean mIsNaturalScrollingEnabled;
 
@@ -756,7 +757,7 @@
     protected void onHandleConfigurationChanged() {
         Trace.beginSection("Launcher#onHandleconfigurationChanged");
         try {
-            if (!initDeviceProfile(mDeviceProfile.inv)) {
+            if (!initDeviceProfile(mDeviceProfile.inv) && !mForceConfigUpdate) {
                 return;
             }
 
@@ -770,6 +771,7 @@
             mModel.rebindCallbacks();
             updateDisallowBack();
         } finally {
+            mForceConfigUpdate = false;
             Trace.endSection();
         }
     }
@@ -1872,13 +1874,18 @@
             }
         }
 
-        // Exit spring loaded mode if necessary after adding the widget
-        Runnable onComplete = MULTI_SELECT_EDIT_MODE.get() ? null
-                : () -> mStateManager.goToState(NORMAL, SPRING_LOADED_EXIT_DELAY);
+        // Exit spring loaded mode if necessary after adding the widget; unless config activity was
+        // started.
+        Runnable onComplete = MULTI_SELECT_EDIT_MODE.get() ? null : () -> mStateManager.goToState(
+                NORMAL, SPRING_LOADED_EXIT_DELAY);
         completeAddAppWidget(appWidgetId, info, boundWidget,
                 addFlowHandler.getProviderInfo(this), addFlowHandler.needsConfigure(),
                 false, widgetPreviewBitmap);
-        mWorkspace.removeExtraEmptyScreenDelayed(delay, false, onComplete);
+        // Remove extra screen if widget drop concluded. If a config activity was started, extra
+        // screen will be removed when we get back its result.
+        if (!isActivityStarted) {
+            mWorkspace.removeExtraEmptyScreenDelayed(delay, false, onComplete);
+        }
     }
 
     public void addPendingItem(PendingAddItemInfo info, int container, int screenId,
@@ -3146,6 +3153,13 @@
         return mAnimationCoordinator;
     }
 
+    /**
+     * Set to force config update when set to true next time onHandleConfigurationChanged is called.
+     */
+    public void setForceConfigUpdate(boolean forceConfigUpdate) {
+        mForceConfigUpdate = forceConfigUpdate;
+    }
+
     @Override
     public View.OnLongClickListener getAllAppsItemLongClickListener() {
         return ItemLongClickListener.INSTANCE_ALL_APPS;
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 15641ab..42a28d6 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -176,7 +176,7 @@
                     () -> LauncherPrefs.get(mContext).removeListener(observer, THEMED_ICONS));
 
             InstallSessionTracker installSessionTracker =
-                    InstallSessionHelper.INSTANCE.get(context).registerInstallTracker(mModel);
+                    InstallSessionHelper.INSTANCE.get(context).registerInstallTracker(callbacks);
             mOnTerminateCallback.add(installSessionTracker::unregister);
         });
 
@@ -266,7 +266,7 @@
     }
 
     private class IconObserver
-            implements IconProvider.IconChangeListener, OnSharedPreferenceChangeListener {
+            implements IconProvider.IconChangeListener, LauncherPrefChangeListener {
 
         @Override
         public void onAppIconChanged(String packageName, UserHandle user) {
@@ -288,7 +288,7 @@
         }
 
         @Override
-        public void onSharedPreferenceChanged(SharedPreferences prefs, String key) {
+        public void onPrefChanged(String key) {
             if (Themes.KEY_THEMED_ICONS.equals(key)) {
                 mIconProvider.setIconThemeSupported(Themes.isThemedIconEnabled(mContext));
                 verifyIconChanged();
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
deleted file mode 100644
index 7ad17d9..0000000
--- a/src/com/android/launcher3/LauncherModel.java
+++ /dev/null
@@ -1,692 +0,0 @@
-/*
- * Copyright (C) 2008 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3;
-
-import static android.app.admin.DevicePolicyManager.ACTION_DEVICE_POLICY_RESOURCE_UPDATED;
-
-import static com.android.launcher3.LauncherAppState.ACTION_FORCE_ROLOAD;
-import static com.android.launcher3.LauncherPrefs.WORK_EDU_STEP;
-import static com.android.launcher3.config.FeatureFlags.IS_STUDIO_BUILD;
-import static com.android.launcher3.icons.cache.BaseIconCache.EMPTY_CLASS_NAME;
-import static com.android.launcher3.model.PackageUpdatedTask.OP_UPDATE;
-import static com.android.launcher3.pm.UserCache.ACTION_PROFILE_AVAILABLE;
-import static com.android.launcher3.pm.UserCache.ACTION_PROFILE_UNAVAILABLE;
-import static com.android.launcher3.testing.shared.TestProtocol.sDebugTracing;
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-
-import android.content.ComponentName;
-import android.content.Context;
-import android.content.Intent;
-import android.content.pm.PackageInstaller;
-import android.content.pm.ShortcutInfo;
-import android.os.UserHandle;
-import android.text.TextUtils;
-import android.util.Log;
-import android.util.Pair;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-import androidx.annotation.WorkerThread;
-
-import com.android.launcher3.celllayout.CellPosMapper;
-import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.icons.IconCache;
-import com.android.launcher3.model.AddWorkspaceItemsTask;
-import com.android.launcher3.model.AllAppsList;
-import com.android.launcher3.model.BaseLauncherBinder;
-import com.android.launcher3.model.BgDataModel;
-import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.CacheDataUpdatedTask;
-import com.android.launcher3.model.ItemInstallQueue;
-import com.android.launcher3.model.LoaderTask;
-import com.android.launcher3.model.ModelDbController;
-import com.android.launcher3.model.ModelDelegate;
-import com.android.launcher3.model.ModelLauncherCallbacks;
-import com.android.launcher3.model.ModelTaskController;
-import com.android.launcher3.model.ModelWriter;
-import com.android.launcher3.model.PackageInstallStateChangedTask;
-import com.android.launcher3.model.PackageUpdatedTask;
-import com.android.launcher3.model.ReloadStringCacheTask;
-import com.android.launcher3.model.ShortcutsChangedTask;
-import com.android.launcher3.model.UserLockStateChangedTask;
-import com.android.launcher3.model.data.AppInfo;
-import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.pm.InstallSessionTracker;
-import com.android.launcher3.pm.PackageInstallInfo;
-import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.shortcuts.ShortcutRequest;
-import com.android.launcher3.util.ApplicationInfoWrapper;
-import com.android.launcher3.util.IntSet;
-import com.android.launcher3.util.ItemInfoMatcher;
-import com.android.launcher3.util.PackageManagerHelper;
-import com.android.launcher3.util.PackageUserKey;
-import com.android.launcher3.util.Preconditions;
-
-import java.io.FileDescriptor;
-import java.io.PrintWriter;
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.concurrent.CancellationException;
-import java.util.function.Consumer;
-import java.util.function.Supplier;
-
-/**
- * Maintains in-memory state of the Launcher. It is expected that there should be only one
- * LauncherModel object held in a static. Also provide APIs for updating the database state
- * for the Launcher.
- */
-public class LauncherModel implements InstallSessionTracker.Callback {
-    private static final boolean DEBUG_RECEIVER = false;
-
-    static final String TAG = "Launcher.Model";
-
-    @NonNull
-    private final LauncherAppState mApp;
-    @NonNull
-    private final PackageManagerHelper mPmHelper;
-    @NonNull
-    private final ModelDbController mModelDbController;
-    @NonNull
-    private final Object mLock = new Object();
-    @Nullable
-    private LoaderTask mLoaderTask;
-    private boolean mIsLoaderTaskRunning;
-
-    // only allow this once per reboot to reload work apps
-    private boolean mShouldReloadWorkProfile = true;
-
-    // Indicates whether the current model data is valid or not.
-    // We start off with everything not loaded. After that, we assume that
-    // our monitoring of the package manager provides all updates and we never
-    // need to do a requery. This is only ever touched from the loader thread.
-    private boolean mModelLoaded;
-    private boolean mModelDestroyed = false;
-    public boolean isModelLoaded() {
-        synchronized (mLock) {
-            return mModelLoaded && mLoaderTask == null && !mModelDestroyed;
-        }
-    }
-
-    @NonNull
-    private final ArrayList<Callbacks> mCallbacksList = new ArrayList<>(1);
-
-    // < only access in worker thread >
-    @NonNull
-    private final AllAppsList mBgAllAppsList;
-
-    /**
-     * All the static data should be accessed on the background thread, A lock should be acquired
-     * on this object when accessing any data from this model.
-     */
-    @NonNull
-    private final BgDataModel mBgDataModel = new BgDataModel();
-
-    @NonNull
-    private final ModelDelegate mModelDelegate;
-
-    private int mLastLoadId = -1;
-
-    // Runnable to check if the shortcuts permission has changed.
-    @NonNull
-    private final Runnable mDataValidationCheck = new Runnable() {
-        @Override
-        public void run() {
-            if (mModelLoaded) {
-                mModelDelegate.validateData();
-            }
-        }
-    };
-
-    LauncherModel(@NonNull final Context context, @NonNull final LauncherAppState app,
-            @NonNull final IconCache iconCache, @NonNull final AppFilter appFilter,
-            @NonNull final PackageManagerHelper pmHelper, final boolean isPrimaryInstance) {
-        mApp = app;
-        mPmHelper = pmHelper;
-        mModelDbController = new ModelDbController(context);
-        mBgAllAppsList = new AllAppsList(iconCache, appFilter);
-        mModelDelegate = ModelDelegate.newInstance(context, app, mPmHelper, mBgAllAppsList,
-                mBgDataModel, isPrimaryInstance);
-    }
-
-    @NonNull
-    public ModelDelegate getModelDelegate() {
-        return mModelDelegate;
-    }
-
-    public ModelDbController getModelDbController() {
-        return mModelDbController;
-    }
-
-    public ModelLauncherCallbacks newModelCallbacks() {
-        return new ModelLauncherCallbacks(this::enqueueModelUpdateTask);
-    }
-
-    /**
-     * Adds the provided items to the workspace.
-     */
-    public void addAndBindAddedWorkspaceItems(
-            @NonNull final List<Pair<ItemInfo, Object>> itemList) {
-        for (Callbacks cb : getCallbacks()) {
-            cb.preAddApps();
-        }
-        enqueueModelUpdateTask(new AddWorkspaceItemsTask(itemList));
-    }
-
-    @NonNull
-    public ModelWriter getWriter(final boolean verifyChanges, CellPosMapper cellPosMapper,
-            @Nullable final Callbacks owner) {
-        return new ModelWriter(mApp.getContext(), this, mBgDataModel, verifyChanges, cellPosMapper,
-                owner);
-    }
-
-    /**
-     * Called when the icon for an app changes, outside of package event
-     */
-    @WorkerThread
-    public void onAppIconChanged(@NonNull final String packageName,
-            @NonNull final UserHandle user) {
-        // Update the icon for the calendar package
-        Context context = mApp.getContext();
-        enqueueModelUpdateTask(new PackageUpdatedTask(OP_UPDATE, user, packageName));
-
-        List<ShortcutInfo> pinnedShortcuts = new ShortcutRequest(context, user)
-                .forPackage(packageName).query(ShortcutRequest.PINNED);
-        if (!pinnedShortcuts.isEmpty()) {
-            enqueueModelUpdateTask(new ShortcutsChangedTask(packageName, pinnedShortcuts, user,
-                    false));
-        }
-    }
-
-    /**
-     * Called when the workspace items have drastically changed
-     */
-    public void onWorkspaceUiChanged() {
-        MODEL_EXECUTOR.execute(mModelDelegate::workspaceLoadComplete);
-    }
-
-    /**
-     * Called when the model is destroyed
-     */
-    public void destroy() {
-        mModelDestroyed = true;
-        MODEL_EXECUTOR.execute(mModelDelegate::destroy);
-    }
-
-    public void onBroadcastIntent(@NonNull final Intent intent) {
-        if (DEBUG_RECEIVER || sDebugTracing) Log.d(TAG, "onReceive intent=" + intent);
-        final String action = intent.getAction();
-        if (Intent.ACTION_LOCALE_CHANGED.equals(action)) {
-            // If we have changed locale we need to clear out the labels in all apps/workspace.
-            forceReload();
-        } else if (ACTION_DEVICE_POLICY_RESOURCE_UPDATED.equals(action)) {
-            enqueueModelUpdateTask(new ReloadStringCacheTask(mModelDelegate));
-        } else if (IS_STUDIO_BUILD && ACTION_FORCE_ROLOAD.equals(action)) {
-            for (Callbacks cb : getCallbacks()) {
-                if (cb instanceof Launcher) {
-                    ((Launcher) cb).recreate();
-                }
-            }
-        }
-    }
-
-    /**
-     * Called then there use a user event
-     * @see UserCache#addUserEventListener
-     */
-    public void onUserEvent(UserHandle user, String action) {
-        if (Intent.ACTION_MANAGED_PROFILE_AVAILABLE.equals(action)
-                && mShouldReloadWorkProfile) {
-            mShouldReloadWorkProfile = false;
-            forceReload();
-        } else if (Intent.ACTION_MANAGED_PROFILE_AVAILABLE.equals(action)
-                || Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE.equals(action)) {
-            mShouldReloadWorkProfile = false;
-            enqueueModelUpdateTask(new PackageUpdatedTask(
-                    PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user));
-        } else if (UserCache.ACTION_PROFILE_LOCKED.equals(action)
-                || UserCache.ACTION_PROFILE_UNLOCKED.equals(action)) {
-            enqueueModelUpdateTask(new UserLockStateChangedTask(
-                    user, UserCache.ACTION_PROFILE_UNLOCKED.equals(action)));
-        } else if (UserCache.ACTION_PROFILE_ADDED.equals(action)
-                || UserCache.ACTION_PROFILE_REMOVED.equals(action)) {
-            forceReload();
-        } else if (ACTION_PROFILE_AVAILABLE.equals(action)
-                || ACTION_PROFILE_UNAVAILABLE.equals(action)) {
-            /*
-             * This broadcast is only available when android.os.Flags.allowPrivateProfile() is set.
-             * For Work-profile this broadcast will be sent in addition to
-             * ACTION_MANAGED_PROFILE_AVAILABLE/UNAVAILABLE.
-             * So effectively, this if block only handles the non-work profile case.
-             */
-            enqueueModelUpdateTask(new PackageUpdatedTask(
-                    PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user));
-        }
-        if (Intent.ACTION_MANAGED_PROFILE_REMOVED.equals(action)) {
-            LauncherPrefs.get(mApp.getContext()).put(WORK_EDU_STEP, 0);
-        }
-    }
-
-    /**
-     * Reloads the workspace items from the DB and re-binds the workspace. This should generally
-     * not be called as DB updates are automatically followed by UI update
-     */
-    public void forceReload() {
-        synchronized (mLock) {
-            // Stop any existing loaders first, so they don't set mModelLoaded to true later
-            stopLoader();
-            mModelLoaded = false;
-        }
-
-        // Start the loader if launcher is already running, otherwise the loader will run,
-        // the next time launcher starts
-        if (hasCallbacks()) {
-            startLoader();
-        }
-    }
-
-    /**
-     * Rebinds all existing callbacks with already loaded model
-     */
-    public void rebindCallbacks() {
-        if (hasCallbacks()) {
-            startLoader();
-        }
-    }
-
-    /**
-     * Removes an existing callback
-     */
-    public void removeCallbacks(@NonNull final Callbacks callbacks) {
-        synchronized (mCallbacksList) {
-            Preconditions.assertUIThread();
-            if (mCallbacksList.remove(callbacks)) {
-                if (stopLoader()) {
-                    // Rebind existing callbacks
-                    startLoader();
-                }
-            }
-        }
-    }
-
-    /**
-     * Adds a callbacks to receive model updates
-     * @return true if workspace load was performed synchronously
-     */
-    public boolean addCallbacksAndLoad(@NonNull final Callbacks callbacks) {
-        synchronized (mLock) {
-            addCallbacks(callbacks);
-            return startLoader(new Callbacks[] { callbacks });
-
-        }
-    }
-
-    /**
-     * Adds a callbacks to receive model updates
-     */
-    public void addCallbacks(@NonNull final Callbacks callbacks) {
-        Preconditions.assertUIThread();
-        synchronized (mCallbacksList) {
-            mCallbacksList.add(callbacks);
-        }
-    }
-
-    /**
-     * Starts the loader. Tries to bind {@params synchronousBindPage} synchronously if possible.
-     * @return true if the page could be bound synchronously.
-     */
-    public boolean startLoader() {
-        return startLoader(new Callbacks[0]);
-    }
-
-    private boolean startLoader(@NonNull final Callbacks[] newCallbacks) {
-        // Enable queue before starting loader. It will get disabled in Launcher#finishBindingItems
-        ItemInstallQueue.INSTANCE.get(mApp.getContext())
-                .pauseModelPush(ItemInstallQueue.FLAG_LOADER_RUNNING);
-        synchronized (mLock) {
-            // If there is already one running, tell it to stop.
-            boolean wasRunning = stopLoader();
-            boolean bindDirectly = mModelLoaded && !mIsLoaderTaskRunning;
-            boolean bindAllCallbacks = wasRunning || !bindDirectly || newCallbacks.length == 0;
-            final Callbacks[] callbacksList = bindAllCallbacks ? getCallbacks() : newCallbacks;
-
-            if (callbacksList.length > 0) {
-                // Clear any pending bind-runnables from the synchronized load process.
-                for (Callbacks cb : callbacksList) {
-                    MAIN_EXECUTOR.execute(cb::clearPendingBinds);
-                }
-
-                BaseLauncherBinder launcherBinder = new BaseLauncherBinder(
-                        mApp, mBgDataModel, mBgAllAppsList, callbacksList);
-                if (bindDirectly) {
-                    // Divide the set of loaded items into those that we are binding synchronously,
-                    // and everything else that is to be bound normally (asynchronously).
-                    launcherBinder.bindWorkspace(bindAllCallbacks, /* isBindSync= */ true);
-                    // For now, continue posting the binding of AllApps as there are other
-                    // issues that arise from that.
-                    launcherBinder.bindAllApps();
-                    launcherBinder.bindDeepShortcuts();
-                    launcherBinder.bindWidgets();
-                    if (FeatureFlags.CHANGE_MODEL_DELEGATE_LOADING_ORDER.get()) {
-                        mModelDelegate.bindAllModelExtras(callbacksList);
-                    }
-                    return true;
-                } else {
-                    stopLoader();
-                    mLoaderTask = new LoaderTask(
-                            mApp, mBgAllAppsList, mBgDataModel, mModelDelegate, launcherBinder);
-
-                    // Always post the loader task, instead of running directly
-                    // (even on same thread) so that we exit any nested synchronized blocks
-                    MODEL_EXECUTOR.post(mLoaderTask);
-                }
-            }
-        }
-        return false;
-    }
-
-    /**
-     * If there is already a loader task running, tell it to stop.
-     * @return true if an existing loader was stopped.
-     */
-    private boolean stopLoader() {
-        synchronized (mLock) {
-            LoaderTask oldTask = mLoaderTask;
-            mLoaderTask = null;
-            if (oldTask != null) {
-                oldTask.stopLocked();
-                return true;
-            }
-            return false;
-        }
-    }
-
-    /**
-     * Loads the model if not loaded
-     * @param callback called with the data model upon successful load or null on model thread.
-     */
-    public void loadAsync(@NonNull final Consumer<BgDataModel> callback) {
-        synchronized (mLock) {
-            if (!mModelLoaded && !mIsLoaderTaskRunning) {
-                startLoader();
-            }
-        }
-        MODEL_EXECUTOR.post(() -> callback.accept(isModelLoaded() ? mBgDataModel : null));
-    }
-
-    @Override
-    public void onInstallSessionCreated(@NonNull final PackageInstallInfo sessionInfo) {
-        if (FeatureFlags.PROMISE_APPS_IN_ALL_APPS.get()) {
-            enqueueModelUpdateTask((taskController, dataModel, apps) -> {
-                apps.addPromiseApp(mApp.getContext(), sessionInfo);
-                taskController.bindApplicationsIfNeeded();
-            });
-        }
-    }
-
-    @Override
-    public void onSessionFailure(@NonNull final String packageName,
-            @NonNull final UserHandle user) {
-        enqueueModelUpdateTask((taskController, dataModel, apps) -> {
-            IconCache iconCache = mApp.getIconCache();
-            final IntSet removedIds = new IntSet();
-            HashSet<WorkspaceItemInfo> archivedWorkspaceItemsToCacheRefresh = new HashSet<>();
-            boolean isAppArchived =
-                    new ApplicationInfoWrapper(mApp.getContext(), packageName, user).isArchived();
-            synchronized (dataModel) {
-                if (isAppArchived) {
-                    // Remove package icon cache entry for archived app in case of a session
-                    // failure.
-                    mApp.getIconCache().remove(
-                            new ComponentName(packageName, packageName + EMPTY_CLASS_NAME),
-                            user);
-                }
-
-                for (ItemInfo info : dataModel.itemsIdMap) {
-                    if (info instanceof WorkspaceItemInfo
-                            && ((WorkspaceItemInfo) info).hasPromiseIconUi()
-                            && user.equals(info.user)
-                            && info.getIntent() != null) {
-                        if (TextUtils.equals(packageName, info.getIntent().getPackage())) {
-                            removedIds.add(info.id);
-                        }
-                        if (((WorkspaceItemInfo) info).isArchived()) {
-                            WorkspaceItemInfo workspaceItem = (WorkspaceItemInfo) info;
-                            // Refresh icons on the workspace for archived apps.
-                            iconCache.getTitleAndIcon(workspaceItem,
-                                    workspaceItem.usingLowResIcon());
-                            archivedWorkspaceItemsToCacheRefresh.add(workspaceItem);
-                        }
-                    }
-                }
-
-                if (isAppArchived) {
-                    apps.updateIconsAndLabels(new HashSet<>(List.of(packageName)), user);
-                }
-            }
-
-            if (!removedIds.isEmpty() && !isAppArchived) {
-                taskController.deleteAndBindComponentsRemoved(
-                        ItemInfoMatcher.ofItemIds(removedIds),
-                        "removed because install session failed");
-            }
-            if (!archivedWorkspaceItemsToCacheRefresh.isEmpty()) {
-                taskController.bindUpdatedWorkspaceItems(
-                        archivedWorkspaceItemsToCacheRefresh.stream().toList());
-            }
-            if (isAppArchived) {
-                taskController.bindApplicationsIfNeeded();
-            }
-        });
-    }
-
-    @Override
-    public void onPackageStateChanged(@NonNull final PackageInstallInfo installInfo) {
-        enqueueModelUpdateTask(new PackageInstallStateChangedTask(installInfo));
-    }
-
-    /**
-     * Updates the icons and label of all pending icons for the provided package name.
-     */
-    @Override
-    public void onUpdateSessionDisplay(@NonNull final PackageUserKey key,
-            @NonNull final PackageInstaller.SessionInfo info) {
-        mApp.getIconCache().updateSessionCache(key, info);
-
-        HashSet<String> packages = new HashSet<>();
-        packages.add(key.mPackageName);
-        enqueueModelUpdateTask(new CacheDataUpdatedTask(
-                CacheDataUpdatedTask.OP_SESSION_UPDATE, key.mUser, packages));
-    }
-
-    public class LoaderTransaction implements AutoCloseable {
-
-        @NonNull
-        private final LoaderTask mTask;
-
-        private LoaderTransaction(@NonNull final LoaderTask task) throws CancellationException {
-            synchronized (mLock) {
-                if (mLoaderTask != task) {
-                    throw new CancellationException("Loader already stopped");
-                }
-                mLastLoadId++;
-                mTask = task;
-                mIsLoaderTaskRunning = true;
-                mModelLoaded = false;
-            }
-        }
-
-        public void commit() {
-            synchronized (mLock) {
-                // Everything loaded bind the data.
-                mModelLoaded = true;
-            }
-        }
-
-        @Override
-        public void close() {
-            synchronized (mLock) {
-                // If we are still the last one to be scheduled, remove ourselves.
-                if (mLoaderTask == mTask) {
-                    mLoaderTask = null;
-                }
-                mIsLoaderTaskRunning = false;
-            }
-        }
-    }
-
-    public LoaderTransaction beginLoader(@NonNull final LoaderTask task)
-            throws CancellationException {
-        return new LoaderTransaction(task);
-    }
-
-    /**
-     * Refreshes the cached shortcuts if the shortcut permission has changed.
-     * Current implementation simply reloads the workspace, but it can be optimized to
-     * use partial updates similar to {@link UserCache}
-     */
-    public void validateModelDataOnResume() {
-        MODEL_EXECUTOR.getHandler().removeCallbacks(mDataValidationCheck);
-        MODEL_EXECUTOR.post(mDataValidationCheck);
-    }
-
-    /**
-     * Called when the icons for packages have been updated in the icon cache.
-     */
-    public void onPackageIconsUpdated(@NonNull final HashSet<String> updatedPackages,
-            @NonNull final UserHandle user) {
-        // If any package icon has changed (app was updated while launcher was dead),
-        // update the corresponding shortcuts.
-        enqueueModelUpdateTask(new CacheDataUpdatedTask(
-                CacheDataUpdatedTask.OP_CACHE_UPDATE, user, updatedPackages));
-    }
-
-    /**
-     * Called when the labels for the widgets has updated in the icon cache.
-     */
-    public void onWidgetLabelsUpdated(@NonNull final HashSet<String> updatedPackages,
-            @NonNull final UserHandle user) {
-        enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
-            dataModel.widgetsModel.onPackageIconsUpdated(updatedPackages, user, mApp);
-            taskController.bindUpdatedWidgets(dataModel);
-        });
-    }
-
-    public void enqueueModelUpdateTask(@NonNull final ModelUpdateTask task) {
-        if (mModelDestroyed) {
-            return;
-        }
-        MODEL_EXECUTOR.execute(() -> {
-            if (!isModelLoaded()) {
-                // Loader has not yet run.
-                return;
-            }
-            ModelTaskController controller = new ModelTaskController(
-                    mApp, mBgDataModel, mBgAllAppsList, this, MAIN_EXECUTOR);
-            task.execute(controller, mBgDataModel, mBgAllAppsList);
-        });
-    }
-
-    /**
-     * A task to be executed on the current callbacks on the UI thread.
-     * If there is no current callbacks, the task is ignored.
-     */
-    public interface CallbackTask {
-
-        void execute(@NonNull Callbacks callbacks);
-    }
-
-    public interface ModelUpdateTask {
-
-        void execute(@NonNull ModelTaskController taskController,
-                @NonNull BgDataModel dataModel, @NonNull AllAppsList apps);
-    }
-
-    public void updateAndBindWorkspaceItem(@NonNull final WorkspaceItemInfo si,
-            @NonNull final ShortcutInfo info) {
-        updateAndBindWorkspaceItem(() -> {
-            si.updateFromDeepShortcutInfo(info, mApp.getContext());
-            mApp.getIconCache().getShortcutIcon(si, info);
-            return si;
-        });
-    }
-
-    /**
-     * Utility method to update a shortcut on the background thread.
-     */
-    public void updateAndBindWorkspaceItem(
-            @NonNull final Supplier<WorkspaceItemInfo> itemProvider) {
-        enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
-            WorkspaceItemInfo info = itemProvider.get();
-            taskController.getModelWriter().updateItemInDatabase(info);
-            ArrayList<WorkspaceItemInfo> update = new ArrayList<>();
-            update.add(info);
-            taskController.bindUpdatedWorkspaceItems(update);
-        });
-    }
-
-    public void refreshAndBindWidgetsAndShortcuts(@Nullable final PackageUserKey packageUser) {
-        enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
-            dataModel.widgetsModel.update(taskController.getApp(), packageUser);
-            taskController.bindUpdatedWidgets(dataModel);
-        });
-    }
-
-    public void dumpState(@Nullable final String prefix, @Nullable final FileDescriptor fd,
-            @NonNull final PrintWriter writer, @NonNull final String[] args) {
-        if (args.length > 0 && TextUtils.equals(args[0], "--all")) {
-            writer.println(prefix + "All apps list: size=" + mBgAllAppsList.data.size());
-            for (AppInfo info : mBgAllAppsList.data) {
-                writer.println(prefix + "   title=\"" + info.title
-                        + "\" bitmapIcon=" + info.bitmap.icon
-                        + " componentName=" + info.componentName.getPackageName());
-            }
-            writer.println();
-        }
-        mModelDelegate.dump(prefix, fd, writer, args);
-        mBgDataModel.dump(prefix, fd, writer, args);
-    }
-
-    /**
-     * Returns true if there are any callbacks attached to the model
-     */
-    public boolean hasCallbacks() {
-        synchronized (mCallbacksList) {
-            return !mCallbacksList.isEmpty();
-        }
-    }
-
-    /**
-     * Returns an array of currently attached callbacks
-     */
-    @NonNull
-    public Callbacks[] getCallbacks() {
-        synchronized (mCallbacksList) {
-            return mCallbacksList.toArray(new Callbacks[mCallbacksList.size()]);
-        }
-    }
-
-    /**
-     * Returns the ID for the last model load. If the load ID doesn't match for a transaction, the
-     * transaction should be ignored.
-     */
-    public int getLastLoadId() {
-        return mLastLoadId;
-    }
-}
diff --git a/src/com/android/launcher3/LauncherModel.kt b/src/com/android/launcher3/LauncherModel.kt
new file mode 100644
index 0000000..a013eaa
--- /dev/null
+++ b/src/com/android/launcher3/LauncherModel.kt
@@ -0,0 +1,499 @@
+/*
+ * Copyright (C) 2008 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3
+
+import android.app.admin.DevicePolicyManager
+import android.content.Context
+import android.content.Intent
+import android.content.pm.ShortcutInfo
+import android.os.UserHandle
+import android.text.TextUtils
+import android.util.Log
+import android.util.Pair
+import androidx.annotation.WorkerThread
+import com.android.launcher3.celllayout.CellPosMapper
+import com.android.launcher3.config.FeatureFlags
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.AddWorkspaceItemsTask
+import com.android.launcher3.model.AllAppsList
+import com.android.launcher3.model.BaseLauncherBinder
+import com.android.launcher3.model.BgDataModel
+import com.android.launcher3.model.CacheDataUpdatedTask
+import com.android.launcher3.model.ItemInstallQueue
+import com.android.launcher3.model.LoaderTask
+import com.android.launcher3.model.ModelDbController
+import com.android.launcher3.model.ModelDelegate
+import com.android.launcher3.model.ModelLauncherCallbacks
+import com.android.launcher3.model.ModelTaskController
+import com.android.launcher3.model.ModelWriter
+import com.android.launcher3.model.PackageUpdatedTask
+import com.android.launcher3.model.ReloadStringCacheTask
+import com.android.launcher3.model.ShortcutsChangedTask
+import com.android.launcher3.model.UserLockStateChangedTask
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.pm.UserCache
+import com.android.launcher3.shortcuts.ShortcutRequest
+import com.android.launcher3.testing.shared.TestProtocol.sDebugTracing
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import com.android.launcher3.util.Executors.MODEL_EXECUTOR
+import com.android.launcher3.util.PackageManagerHelper
+import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.Preconditions
+import java.io.FileDescriptor
+import java.io.PrintWriter
+import java.util.concurrent.CancellationException
+import java.util.function.Consumer
+
+/**
+ * Maintains in-memory state of the Launcher. It is expected that there should be only one
+ * LauncherModel object held in a static. Also provide APIs for updating the database state for the
+ * Launcher.
+ */
+class LauncherModel(
+    private val context: Context,
+    private val mApp: LauncherAppState,
+    private val iconCache: IconCache,
+    private val appFilter: AppFilter,
+    private val mPmHelper: PackageManagerHelper,
+    isPrimaryInstance: Boolean,
+) {
+
+    private val mCallbacksList = ArrayList<BgDataModel.Callbacks>(1)
+
+    // < only access in worker thread >
+    private val mBgAllAppsList = AllAppsList(iconCache, appFilter)
+
+    /**
+     * All the static data should be accessed on the background thread, A lock should be acquired on
+     * this object when accessing any data from this model.
+     */
+    private val mBgDataModel = BgDataModel()
+
+    val modelDelegate: ModelDelegate =
+        ModelDelegate.newInstance(
+            context,
+            mApp,
+            mPmHelper,
+            mBgAllAppsList,
+            mBgDataModel,
+            isPrimaryInstance,
+        )
+
+    val modelDbController = ModelDbController(context)
+
+    private val mLock = Any()
+
+    private var mLoaderTask: LoaderTask? = null
+    private var mIsLoaderTaskRunning = false
+
+    // only allow this once per reboot to reload work apps
+    private var mShouldReloadWorkProfile = true
+
+    // Indicates whether the current model data is valid or not.
+    // We start off with everything not loaded. After that, we assume that
+    // our monitoring of the package manager provides all updates and we never
+    // need to do a requery. This is only ever touched from the loader thread.
+    private var mModelLoaded = false
+    private var mModelDestroyed = false
+
+    fun isModelLoaded() =
+        synchronized(mLock) { mModelLoaded && mLoaderTask == null && !mModelDestroyed }
+
+    /**
+     * Returns the ID for the last model load. If the load ID doesn't match for a transaction, the
+     * transaction should be ignored.
+     */
+    var lastLoadId: Int = -1
+        private set
+
+    // Runnable to check if the shortcuts permission has changed.
+    private val mDataValidationCheck = Runnable {
+        if (mModelLoaded) {
+            modelDelegate.validateData()
+        }
+    }
+
+    fun newModelCallbacks() = ModelLauncherCallbacks(this::enqueueModelUpdateTask)
+
+    /** Adds the provided items to the workspace. */
+    fun addAndBindAddedWorkspaceItems(itemList: List<Pair<ItemInfo?, Any?>?>) {
+        callbacks.forEach { it.preAddApps() }
+        enqueueModelUpdateTask(AddWorkspaceItemsTask(itemList))
+    }
+
+    fun getWriter(
+        verifyChanges: Boolean,
+        cellPosMapper: CellPosMapper?,
+        owner: BgDataModel.Callbacks?,
+    ) = ModelWriter(mApp.context, this, mBgDataModel, verifyChanges, cellPosMapper, owner)
+
+    /** Called when the icon for an app changes, outside of package event */
+    @WorkerThread
+    fun onAppIconChanged(packageName: String, user: UserHandle) {
+        // Update the icon for the calendar package
+        enqueueModelUpdateTask(PackageUpdatedTask(PackageUpdatedTask.OP_UPDATE, user, packageName))
+        ShortcutRequest(context, user).forPackage(packageName).query(ShortcutRequest.PINNED).let {
+            if (it.isNotEmpty()) {
+                enqueueModelUpdateTask(ShortcutsChangedTask(packageName, it, user, false))
+            }
+        }
+    }
+
+    /** Called when the workspace items have drastically changed */
+    fun onWorkspaceUiChanged() {
+        MODEL_EXECUTOR.execute(modelDelegate::workspaceLoadComplete)
+    }
+
+    /** Called when the model is destroyed */
+    fun destroy() {
+        mModelDestroyed = true
+        MODEL_EXECUTOR.execute(modelDelegate::destroy)
+    }
+
+    fun onBroadcastIntent(intent: Intent) {
+        if (DEBUG_RECEIVER || sDebugTracing) Log.d(TAG, "onReceive intent=$intent")
+        when (intent.action) {
+            Intent.ACTION_LOCALE_CHANGED,
+            LauncherAppState.ACTION_FORCE_ROLOAD ->
+                // If we have changed locale we need to clear out the labels in all apps/workspace.
+                forceReload()
+            DevicePolicyManager.ACTION_DEVICE_POLICY_RESOURCE_UPDATED ->
+                enqueueModelUpdateTask(ReloadStringCacheTask(this.modelDelegate))
+        }
+    }
+
+    /**
+     * Called then there use a user event
+     *
+     * @see UserCache.addUserEventListener
+     */
+    fun onUserEvent(user: UserHandle, action: String) {
+        when (action) {
+            Intent.ACTION_MANAGED_PROFILE_AVAILABLE -> {
+                if (mShouldReloadWorkProfile) {
+                    forceReload()
+                } else {
+                    enqueueModelUpdateTask(
+                        PackageUpdatedTask(PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user)
+                    )
+                }
+                mShouldReloadWorkProfile = false
+            }
+            Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE -> {
+                mShouldReloadWorkProfile = false
+                enqueueModelUpdateTask(
+                    PackageUpdatedTask(PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user)
+                )
+            }
+            UserCache.ACTION_PROFILE_LOCKED ->
+                enqueueModelUpdateTask(UserLockStateChangedTask(user, false))
+            UserCache.ACTION_PROFILE_UNLOCKED ->
+                enqueueModelUpdateTask(UserLockStateChangedTask(user, true))
+            Intent.ACTION_MANAGED_PROFILE_REMOVED -> {
+                LauncherPrefs.get(mApp.context).put(LauncherPrefs.WORK_EDU_STEP, 0)
+                forceReload()
+            }
+            UserCache.ACTION_PROFILE_ADDED,
+            UserCache.ACTION_PROFILE_REMOVED -> forceReload()
+            UserCache.ACTION_PROFILE_AVAILABLE,
+            UserCache.ACTION_PROFILE_UNAVAILABLE -> {
+                // This broadcast is only available when android.os.Flags.allowPrivateProfile() is
+                // set. For Work-profile this broadcast will be sent in addition to
+                // ACTION_MANAGED_PROFILE_AVAILABLE/UNAVAILABLE. So effectively, this if block only
+                // handles the non-work profile case.
+                enqueueModelUpdateTask(
+                    PackageUpdatedTask(PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user)
+                )
+            }
+        }
+    }
+
+    /**
+     * Reloads the workspace items from the DB and re-binds the workspace. This should generally not
+     * be called as DB updates are automatically followed by UI update
+     */
+    fun forceReload() {
+        synchronized(mLock) {
+            // Stop any existing loaders first, so they don't set mModelLoaded to true later
+            stopLoader()
+            mModelLoaded = false
+        }
+        rebindCallbacks()
+    }
+
+    /** Rebinds all existing callbacks with already loaded model */
+    fun rebindCallbacks() {
+        if (hasCallbacks()) {
+            startLoader()
+        }
+    }
+
+    /** Removes an existing callback */
+    fun removeCallbacks(callbacks: BgDataModel.Callbacks) {
+        synchronized(mCallbacksList) {
+            Preconditions.assertUIThread()
+            if (mCallbacksList.remove(callbacks)) {
+                if (stopLoader()) {
+                    // Rebind existing callbacks
+                    startLoader()
+                }
+            }
+        }
+    }
+
+    /**
+     * Adds a callbacks to receive model updates
+     *
+     * @return true if workspace load was performed synchronously
+     */
+    fun addCallbacksAndLoad(callbacks: BgDataModel.Callbacks): Boolean {
+        synchronized(mLock) {
+            addCallbacks(callbacks)
+            return startLoader(arrayOf(callbacks))
+        }
+    }
+
+    /** Adds a callbacks to receive model updates */
+    fun addCallbacks(callbacks: BgDataModel.Callbacks) {
+        Preconditions.assertUIThread()
+        synchronized(mCallbacksList) { mCallbacksList.add(callbacks) }
+    }
+
+    /**
+     * Starts the loader. Tries to bind {@params synchronousBindPage} synchronously if possible.
+     *
+     * @return true if the page could be bound synchronously.
+     */
+    fun startLoader() = startLoader(arrayOf())
+
+    private fun startLoader(newCallbacks: Array<BgDataModel.Callbacks>): Boolean {
+        // Enable queue before starting loader. It will get disabled in Launcher#finishBindingItems
+        ItemInstallQueue.INSTANCE.get(context).pauseModelPush(ItemInstallQueue.FLAG_LOADER_RUNNING)
+        synchronized(mLock) {
+            // If there is already one running, tell it to stop.
+            val wasRunning = stopLoader()
+            val bindDirectly = mModelLoaded && !mIsLoaderTaskRunning
+            val bindAllCallbacks = wasRunning || !bindDirectly || newCallbacks.isEmpty()
+            val callbacksList = if (bindAllCallbacks) callbacks else newCallbacks
+            if (callbacksList.isNotEmpty()) {
+                // Clear any pending bind-runnables from the synchronized load process.
+                callbacksList.forEach { MAIN_EXECUTOR.execute(it::clearPendingBinds) }
+
+                val launcherBinder =
+                    BaseLauncherBinder(mApp, mBgDataModel, mBgAllAppsList, callbacksList)
+                if (bindDirectly) {
+                    // Divide the set of loaded items into those that we are binding synchronously,
+                    // and everything else that is to be bound normally (asynchronously).
+                    launcherBinder.bindWorkspace(bindAllCallbacks, /* isBindSync= */ true)
+                    // For now, continue posting the binding of AllApps as there are other
+                    // issues that arise from that.
+                    launcherBinder.bindAllApps()
+                    launcherBinder.bindDeepShortcuts()
+                    launcherBinder.bindWidgets()
+                    if (FeatureFlags.CHANGE_MODEL_DELEGATE_LOADING_ORDER.get()) {
+                        this.modelDelegate.bindAllModelExtras(callbacksList)
+                    }
+                    return true
+                } else {
+                    mLoaderTask =
+                        LoaderTask(
+                            mApp,
+                            mBgAllAppsList,
+                            mBgDataModel,
+                            this.modelDelegate,
+                            launcherBinder,
+                        )
+
+                    // Always post the loader task, instead of running directly
+                    // (even on same thread) so that we exit any nested synchronized blocks
+                    MODEL_EXECUTOR.post(mLoaderTask)
+                }
+            }
+        }
+        return false
+    }
+
+    /**
+     * If there is already a loader task running, tell it to stop.
+     *
+     * @return true if an existing loader was stopped.
+     */
+    private fun stopLoader(): Boolean {
+        synchronized(mLock) {
+            val oldTask: LoaderTask? = mLoaderTask
+            mLoaderTask = null
+            if (oldTask != null) {
+                oldTask.stopLocked()
+                return true
+            }
+            return false
+        }
+    }
+
+    /**
+     * Loads the model if not loaded
+     *
+     * @param callback called with the data model upon successful load or null on model thread.
+     */
+    fun loadAsync(callback: Consumer<BgDataModel?>) {
+        synchronized(mLock) {
+            if (!mModelLoaded && !mIsLoaderTaskRunning) {
+                startLoader()
+            }
+        }
+        MODEL_EXECUTOR.post { callback.accept(if (isModelLoaded()) mBgDataModel else null) }
+    }
+
+    inner class LoaderTransaction(task: LoaderTask) : AutoCloseable {
+        private var mTask: LoaderTask? = null
+
+        init {
+            synchronized(mLock) {
+                if (mLoaderTask !== task) {
+                    throw CancellationException("Loader already stopped")
+                }
+                this@LauncherModel.lastLoadId++
+                mTask = task
+                mIsLoaderTaskRunning = true
+                mModelLoaded = false
+            }
+        }
+
+        fun commit() {
+            synchronized(mLock) {
+                // Everything loaded bind the data.
+                mModelLoaded = true
+            }
+        }
+
+        override fun close() {
+            synchronized(mLock) {
+                // If we are still the last one to be scheduled, remove ourselves.
+                if (mLoaderTask === mTask) {
+                    mLoaderTask = null
+                }
+                mIsLoaderTaskRunning = false
+            }
+        }
+    }
+
+    @Throws(CancellationException::class)
+    fun beginLoader(task: LoaderTask) = LoaderTransaction(task)
+
+    /**
+     * Refreshes the cached shortcuts if the shortcut permission has changed. Current implementation
+     * simply reloads the workspace, but it can be optimized to use partial updates similar to
+     * [UserCache]
+     */
+    fun validateModelDataOnResume() {
+        MODEL_EXECUTOR.handler.removeCallbacks(mDataValidationCheck)
+        MODEL_EXECUTOR.post(mDataValidationCheck)
+    }
+
+    /** Called when the icons for packages have been updated in the icon cache. */
+    fun onPackageIconsUpdated(updatedPackages: HashSet<String?>, user: UserHandle) {
+        // If any package icon has changed (app was updated while launcher was dead),
+        // update the corresponding shortcuts.
+        enqueueModelUpdateTask(
+            CacheDataUpdatedTask(CacheDataUpdatedTask.OP_CACHE_UPDATE, user, updatedPackages)
+        )
+    }
+
+    /** Called when the labels for the widgets has updated in the icon cache. */
+    fun onWidgetLabelsUpdated(updatedPackages: HashSet<String?>, user: UserHandle) {
+        enqueueModelUpdateTask { taskController, dataModel, _ ->
+            dataModel.widgetsModel.onPackageIconsUpdated(updatedPackages, user, mApp)
+            taskController.bindUpdatedWidgets(dataModel)
+        }
+    }
+
+    fun enqueueModelUpdateTask(task: ModelUpdateTask) {
+        if (mModelDestroyed) {
+            return
+        }
+        MODEL_EXECUTOR.execute {
+            if (!isModelLoaded()) {
+                // Loader has not yet run.
+                return@execute
+            }
+            task.execute(
+                ModelTaskController(mApp, mBgDataModel, mBgAllAppsList, this, MAIN_EXECUTOR),
+                mBgDataModel,
+                mBgAllAppsList,
+            )
+        }
+    }
+
+    /**
+     * A task to be executed on the current callbacks on the UI thread. If there is no current
+     * callbacks, the task is ignored.
+     */
+    fun interface CallbackTask {
+        fun execute(callbacks: BgDataModel.Callbacks)
+    }
+
+    fun interface ModelUpdateTask {
+        fun execute(taskController: ModelTaskController, dataModel: BgDataModel, apps: AllAppsList)
+    }
+
+    fun updateAndBindWorkspaceItem(si: WorkspaceItemInfo, info: ShortcutInfo) {
+        enqueueModelUpdateTask { taskController, _, _ ->
+            si.updateFromDeepShortcutInfo(info, context)
+            iconCache.getShortcutIcon(si, info)
+            taskController.getModelWriter().updateItemInDatabase(si)
+            taskController.bindUpdatedWorkspaceItems(listOf(si))
+        }
+    }
+
+    fun refreshAndBindWidgetsAndShortcuts(packageUser: PackageUserKey?) {
+        enqueueModelUpdateTask { taskController, dataModel, _ ->
+            dataModel.widgetsModel.update(taskController.app, packageUser)
+            taskController.bindUpdatedWidgets(dataModel)
+        }
+    }
+
+    fun dumpState(prefix: String?, fd: FileDescriptor?, writer: PrintWriter, args: Array<String?>) {
+        if (args.isNotEmpty() && TextUtils.equals(args[0], "--all")) {
+            writer.println(prefix + "All apps list: size=" + mBgAllAppsList.data.size)
+            for (info in mBgAllAppsList.data) {
+                writer.println(
+                    "$prefix   title=\"${info.title}\" bitmapIcon=${info.bitmap.icon} componentName=${info.targetPackage}"
+                )
+            }
+            writer.println()
+        }
+        modelDelegate.dump(prefix, fd, writer, args)
+        mBgDataModel.dump(prefix, fd, writer, args)
+    }
+
+    /** Returns true if there are any callbacks attached to the model */
+    fun hasCallbacks() = synchronized(mCallbacksList) { mCallbacksList.isNotEmpty() }
+
+    /** Returns an array of currently attached callbacks */
+    val callbacks: Array<BgDataModel.Callbacks>
+        get() {
+            synchronized(mCallbacksList) {
+                return mCallbacksList.toTypedArray<BgDataModel.Callbacks>()
+            }
+        }
+
+    companion object {
+        private const val DEBUG_RECEIVER = false
+
+        const val TAG = "Launcher.Model"
+    }
+}
diff --git a/src/com/android/launcher3/LauncherPrefChangeListener.java b/src/com/android/launcher3/LauncherPrefChangeListener.java
new file mode 100644
index 0000000..3e9a846
--- /dev/null
+++ b/src/com/android/launcher3/LauncherPrefChangeListener.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3;
+
+import android.content.SharedPreferences;
+import android.content.SharedPreferences.OnSharedPreferenceChangeListener;
+
+/**
+ * Listener for changes in [LauncherPrefs].
+ * <p>
+ * The listener also serves as an [OnSharedPreferenceChangeListener] where
+ * [onSharedPreferenceChanged] delegates to [onPrefChanged]. Overriding [onSharedPreferenceChanged]
+ * breaks compatibility with [SharedPreferences].
+ */
+public interface LauncherPrefChangeListener extends OnSharedPreferenceChangeListener {
+
+    /** Callback invoked when the preference for [key] has changed. */
+    void onPrefChanged(String key);
+
+    @Override
+    default void onSharedPreferenceChanged(SharedPreferences sharedPreferences, String key) {
+        onPrefChanged(key);
+    }
+}
diff --git a/src/com/android/launcher3/LauncherPrefs.kt b/src/com/android/launcher3/LauncherPrefs.kt
index 13181e8..7ebfc18 100644
--- a/src/com/android/launcher3/LauncherPrefs.kt
+++ b/src/com/android/launcher3/LauncherPrefs.kt
@@ -18,7 +18,6 @@
 import android.content.Context
 import android.content.Context.MODE_PRIVATE
 import android.content.SharedPreferences
-import android.content.SharedPreferences.OnSharedPreferenceChangeListener
 import androidx.annotation.VisibleForTesting
 import com.android.launcher3.BuildConfig.WIDGET_ON_FIRST_SCREEN
 import com.android.launcher3.LauncherFiles.DEVICE_PREFERENCES_KEY
@@ -34,11 +33,177 @@
 import com.android.launcher3.util.Themes
 
 /**
- * Use same context for shared preferences, so that we use a single cached instance
+ * Manages Launcher [SharedPreferences] through [Item] instances.
  *
  * TODO(b/262721340): Replace all direct SharedPreference refs with LauncherPrefs / Item methods.
  */
-class LauncherPrefs(private val encryptedContext: Context) : SafeCloseable {
+abstract class LauncherPrefs : SafeCloseable {
+
+    /** Returns the value with type [T] for [item]. */
+    abstract fun <T> get(item: ContextualItem<T>): T
+
+    /** Returns the value with type [T] for [item]. */
+    abstract fun <T> get(item: ConstantItem<T>): T
+
+    /** Stores the values for each item in preferences. */
+    abstract fun put(vararg itemsToValues: Pair<Item, Any>)
+
+    /** Stores the [value] with type [T] for [item] in preferences. */
+    abstract fun <T : Any> put(item: Item, value: T)
+
+    /** Synchronous version of [put]. */
+    abstract fun putSync(vararg itemsToValues: Pair<Item, Any>)
+
+    /** Registers [listener] for [items]. */
+    abstract fun addListener(listener: LauncherPrefChangeListener, vararg items: Item)
+
+    /** Unregisters [listener] for [items]. */
+    abstract fun removeListener(listener: LauncherPrefChangeListener, vararg items: Item)
+
+    /** Returns `true` iff all [items] have a value. */
+    abstract fun has(vararg items: Item): Boolean
+
+    /** Removes the value for each item in [items]. */
+    abstract fun remove(vararg items: Item)
+
+    /** Synchronous version of [remove]. */
+    abstract fun removeSync(vararg items: Item)
+
+    companion object {
+        @VisibleForTesting const val BOOT_AWARE_PREFS_KEY = "boot_aware_prefs"
+
+        @JvmField
+        var INSTANCE = MainThreadInitializedObject<LauncherPrefs> { LauncherPrefsImpl(it) }
+
+        @JvmStatic fun get(context: Context): LauncherPrefs = INSTANCE.get(context)
+
+        const val TASKBAR_PINNING_KEY = "TASKBAR_PINNING_KEY"
+        const val TASKBAR_PINNING_DESKTOP_MODE_KEY = "TASKBAR_PINNING_DESKTOP_MODE_KEY"
+        const val SHOULD_SHOW_SMARTSPACE_KEY = "SHOULD_SHOW_SMARTSPACE_KEY"
+        @JvmField
+        val ICON_STATE = nonRestorableItem("pref_icon_shape_path", "", EncryptionType.ENCRYPTED)
+
+        @JvmField
+        val ENABLE_TWOLINE_ALLAPPS_TOGGLE = backedUpItem("pref_enable_two_line_toggle", false)
+        @JvmField
+        val THEMED_ICONS = backedUpItem(Themes.KEY_THEMED_ICONS, false, EncryptionType.ENCRYPTED)
+        @JvmField val PROMISE_ICON_IDS = backedUpItem(InstallSessionHelper.PROMISE_ICON_IDS, "")
+        @JvmField val WORK_EDU_STEP = backedUpItem("showed_work_profile_edu", 0)
+        @JvmField
+        val WORKSPACE_SIZE =
+            backedUpItem(DeviceGridState.KEY_WORKSPACE_SIZE, "", EncryptionType.ENCRYPTED)
+        @JvmField
+        val HOTSEAT_COUNT =
+            backedUpItem(DeviceGridState.KEY_HOTSEAT_COUNT, -1, EncryptionType.ENCRYPTED)
+        @JvmField
+        val TASKBAR_PINNING =
+            backedUpItem(TASKBAR_PINNING_KEY, false, EncryptionType.DEVICE_PROTECTED)
+        @JvmField
+        val TASKBAR_PINNING_IN_DESKTOP_MODE =
+            backedUpItem(TASKBAR_PINNING_DESKTOP_MODE_KEY, true, EncryptionType.DEVICE_PROTECTED)
+
+        @JvmField
+        val DEVICE_TYPE =
+            backedUpItem(
+                DeviceGridState.KEY_DEVICE_TYPE,
+                InvariantDeviceProfile.TYPE_PHONE,
+                EncryptionType.ENCRYPTED,
+            )
+        @JvmField
+        val DB_FILE = backedUpItem(DeviceGridState.KEY_DB_FILE, "", EncryptionType.ENCRYPTED)
+        @JvmField
+        val SHOULD_SHOW_SMARTSPACE =
+            backedUpItem(
+                SHOULD_SHOW_SMARTSPACE_KEY,
+                WIDGET_ON_FIRST_SCREEN,
+                EncryptionType.DEVICE_PROTECTED,
+            )
+        @JvmField
+        val RESTORE_DEVICE =
+            backedUpItem(
+                RestoreDbTask.RESTORED_DEVICE_TYPE,
+                InvariantDeviceProfile.TYPE_PHONE,
+                EncryptionType.ENCRYPTED,
+            )
+        @JvmField
+        val IS_FIRST_LOAD_AFTER_RESTORE =
+            nonRestorableItem(FIRST_LOAD_AFTER_RESTORE_KEY, false, EncryptionType.ENCRYPTED)
+        @JvmField val APP_WIDGET_IDS = backedUpItem(RestoreDbTask.APPWIDGET_IDS, "")
+        @JvmField val OLD_APP_WIDGET_IDS = backedUpItem(RestoreDbTask.APPWIDGET_OLD_IDS, "")
+        @JvmField
+        val GRID_NAME =
+            ConstantItem(
+                "idp_grid_name",
+                isBackedUp = true,
+                defaultValue = null,
+                encryptionType = EncryptionType.ENCRYPTED,
+                type = String::class.java,
+            )
+        @JvmField
+        val ALLOW_ROTATION =
+            backedUpItem(RotationHelper.ALLOW_ROTATION_PREFERENCE_KEY, Boolean::class.java) {
+                RotationHelper.getAllowRotationDefaultValue(DisplayController.INSTANCE.get(it).info)
+            }
+
+        // Preferences for widget configurations
+        @JvmField
+        val RECONFIGURABLE_WIDGET_EDUCATION_TIP_SEEN =
+            backedUpItem("launcher.reconfigurable_widget_education_tip_seen", false)
+
+        @JvmStatic
+        fun <T> backedUpItem(
+            sharedPrefKey: String,
+            defaultValue: T,
+            encryptionType: EncryptionType = EncryptionType.ENCRYPTED,
+        ): ConstantItem<T> =
+            ConstantItem(sharedPrefKey, isBackedUp = true, defaultValue, encryptionType)
+
+        @JvmStatic
+        fun <T> backedUpItem(
+            sharedPrefKey: String,
+            type: Class<out T>,
+            encryptionType: EncryptionType = EncryptionType.ENCRYPTED,
+            defaultValueFromContext: (c: Context) -> T,
+        ): ContextualItem<T> =
+            ContextualItem(
+                sharedPrefKey,
+                isBackedUp = true,
+                defaultValueFromContext,
+                encryptionType,
+                type,
+            )
+
+        @JvmStatic
+        fun <T> nonRestorableItem(
+            sharedPrefKey: String,
+            defaultValue: T,
+            encryptionType: EncryptionType = EncryptionType.ENCRYPTED,
+        ): ConstantItem<T> =
+            ConstantItem(sharedPrefKey, isBackedUp = false, defaultValue, encryptionType)
+
+        @Deprecated("Don't use shared preferences directly. Use other LauncherPref methods.")
+        @JvmStatic
+        fun getPrefs(context: Context): SharedPreferences {
+            // Use application context for shared preferences, so we use single cached instance
+            return context.applicationContext.getSharedPreferences(
+                SHARED_PREFERENCES_KEY,
+                MODE_PRIVATE,
+            )
+        }
+
+        @Deprecated("Don't use shared preferences directly. Use other LauncherPref methods.")
+        @JvmStatic
+        fun getDevicePrefs(context: Context): SharedPreferences {
+            // Use application context for shared preferences, so we use a single cached instance
+            return context.applicationContext.getSharedPreferences(
+                DEVICE_PREFERENCES_KEY,
+                MODE_PRIVATE,
+            )
+        }
+    }
+}
+
+private class LauncherPrefsImpl(private val encryptedContext: Context) : LauncherPrefs() {
     private val deviceProtectedStorageContext =
         encryptedContext.createDeviceProtectedStorageContext()
 
@@ -54,11 +219,11 @@
         else item.encryptedPrefs
 
     /** Wrapper around `getInner` for a `ContextualItem` */
-    fun <T> get(item: ContextualItem<T>): T =
+    override fun <T> get(item: ContextualItem<T>): T =
         getInner(item, item.defaultValueFromContext(encryptedContext))
 
     /** Wrapper around `getInner` for an `Item` */
-    fun <T> get(item: ConstantItem<T>): T = getInner(item, item.defaultValue)
+    override fun <T> get(item: ConstantItem<T>): T = getInner(item, item.defaultValue)
 
     /**
      * Retrieves the value for an [Item] from [SharedPreferences]. It handles method typing via the
@@ -97,17 +262,17 @@
      * prepareToPutValue(itemsToValues) for every distinct `SharedPreferences` file present in the
      * provided item configurations.
      */
-    fun put(vararg itemsToValues: Pair<Item, Any>): Unit =
+    override fun put(vararg itemsToValues: Pair<Item, Any>): Unit =
         prepareToPutValues(itemsToValues).forEach { it.apply() }
 
     /** See referenced `put` method above. */
-    fun <T : Any> put(item: Item, value: T): Unit = put(item.to(value))
+    override fun <T : Any> put(item: Item, value: T): Unit = put(item.to(value))
 
     /**
      * Synchronously stores all the values provided according to their associated Item
      * configuration.
      */
-    fun putSync(vararg itemsToValues: Pair<Item, Any>): Unit =
+    override fun putSync(vararg itemsToValues: Pair<Item, Any>): Unit =
         prepareToPutValues(itemsToValues).forEach { it.commit() }
 
     /**
@@ -152,7 +317,7 @@
     @Suppress("UNCHECKED_CAST")
     private fun SharedPreferences.Editor.putValue(
         item: Item,
-        value: Any?
+        value: Any?,
     ): SharedPreferences.Editor =
         when (item.type) {
             String::class.java -> putString(item.sharedPrefKey, value as? String)
@@ -176,7 +341,7 @@
      * `SharedPreferences` files associated with the provided list of items. The listener will need
      * to filter update notifications so they don't activate for non-relevant updates.
      */
-    fun addListener(listener: OnSharedPreferenceChangeListener, vararg items: Item) {
+    override fun addListener(listener: LauncherPrefChangeListener, vararg items: Item) {
         items
             .map { chooseSharedPreferences(it) }
             .distinct()
@@ -187,7 +352,7 @@
      * Stops the listener from getting notified of any more updates to any of the
      * `SharedPreferences` files associated with any of the provided list of [Item].
      */
-    fun removeListener(listener: OnSharedPreferenceChangeListener, vararg items: Item) {
+    override fun removeListener(listener: LauncherPrefChangeListener, vararg items: Item) {
         // If a listener is not registered to a SharedPreference, unregistering it does nothing
         items
             .map { chooseSharedPreferences(it) }
@@ -199,7 +364,7 @@
      * Checks if all the provided [Item] have values stored in their corresponding
      * `SharedPreferences` files.
      */
-    fun has(vararg items: Item): Boolean {
+    override fun has(vararg items: Item): Boolean {
         items
             .groupBy { chooseSharedPreferences(it) }
             .forEach { (prefs, itemsSublist) ->
@@ -211,10 +376,10 @@
     /**
      * Asynchronously removes the [Item]'s value from its corresponding `SharedPreferences` file.
      */
-    fun remove(vararg items: Item) = prepareToRemove(items).forEach { it.apply() }
+    override fun remove(vararg items: Item) = prepareToRemove(items).forEach { it.apply() }
 
     /** Synchronously removes the [Item]'s value from its corresponding `SharedPreferences` file. */
-    fun removeSync(vararg items: Item) = prepareToRemove(items).forEach { it.commit() }
+    override fun removeSync(vararg items: Item) = prepareToRemove(items).forEach { it.commit() }
 
     /**
      * Removes the key value pairs stored in `SharedPreferences` for each corresponding Item. If the
@@ -244,138 +409,6 @@
     }
 
     override fun close() {}
-
-    companion object {
-        @VisibleForTesting const val BOOT_AWARE_PREFS_KEY = "boot_aware_prefs"
-
-        @JvmField var INSTANCE = MainThreadInitializedObject { LauncherPrefs(it) }
-
-        @JvmStatic fun get(context: Context): LauncherPrefs = INSTANCE.get(context)
-
-        const val TASKBAR_PINNING_KEY = "TASKBAR_PINNING_KEY"
-        const val TASKBAR_PINNING_DESKTOP_MODE_KEY = "TASKBAR_PINNING_DESKTOP_MODE_KEY"
-        const val SHOULD_SHOW_SMARTSPACE_KEY = "SHOULD_SHOW_SMARTSPACE_KEY"
-        @JvmField
-        val ICON_STATE = nonRestorableItem("pref_icon_shape_path", "", EncryptionType.ENCRYPTED)
-
-        @JvmField
-        val ENABLE_TWOLINE_ALLAPPS_TOGGLE = backedUpItem("pref_enable_two_line_toggle", false)
-        @JvmField
-        val THEMED_ICONS = backedUpItem(Themes.KEY_THEMED_ICONS, false, EncryptionType.ENCRYPTED)
-        @JvmField val PROMISE_ICON_IDS = backedUpItem(InstallSessionHelper.PROMISE_ICON_IDS, "")
-        @JvmField val WORK_EDU_STEP = backedUpItem("showed_work_profile_edu", 0)
-        @JvmField
-        val WORKSPACE_SIZE =
-            backedUpItem(DeviceGridState.KEY_WORKSPACE_SIZE, "", EncryptionType.ENCRYPTED)
-        @JvmField
-        val HOTSEAT_COUNT =
-            backedUpItem(DeviceGridState.KEY_HOTSEAT_COUNT, -1, EncryptionType.ENCRYPTED)
-        @JvmField
-        val TASKBAR_PINNING =
-            backedUpItem(TASKBAR_PINNING_KEY, false, EncryptionType.DEVICE_PROTECTED)
-        @JvmField
-        val TASKBAR_PINNING_IN_DESKTOP_MODE =
-            backedUpItem(TASKBAR_PINNING_DESKTOP_MODE_KEY, true, EncryptionType.DEVICE_PROTECTED)
-
-        @JvmField
-        val DEVICE_TYPE =
-            backedUpItem(
-                DeviceGridState.KEY_DEVICE_TYPE,
-                InvariantDeviceProfile.TYPE_PHONE,
-                EncryptionType.ENCRYPTED
-            )
-        @JvmField
-        val DB_FILE = backedUpItem(DeviceGridState.KEY_DB_FILE, "", EncryptionType.ENCRYPTED)
-        @JvmField
-        val SHOULD_SHOW_SMARTSPACE =
-            backedUpItem(
-                SHOULD_SHOW_SMARTSPACE_KEY,
-                WIDGET_ON_FIRST_SCREEN,
-                EncryptionType.DEVICE_PROTECTED
-            )
-        @JvmField
-        val RESTORE_DEVICE =
-            backedUpItem(
-                RestoreDbTask.RESTORED_DEVICE_TYPE,
-                InvariantDeviceProfile.TYPE_PHONE,
-                EncryptionType.ENCRYPTED
-            )
-        @JvmField
-        val IS_FIRST_LOAD_AFTER_RESTORE =
-            nonRestorableItem(FIRST_LOAD_AFTER_RESTORE_KEY, false, EncryptionType.ENCRYPTED)
-        @JvmField val APP_WIDGET_IDS = backedUpItem(RestoreDbTask.APPWIDGET_IDS, "")
-        @JvmField val OLD_APP_WIDGET_IDS = backedUpItem(RestoreDbTask.APPWIDGET_OLD_IDS, "")
-        @JvmField
-        val GRID_NAME =
-            ConstantItem(
-                "idp_grid_name",
-                isBackedUp = true,
-                defaultValue = null,
-                encryptionType = EncryptionType.ENCRYPTED,
-                type = String::class.java
-            )
-        @JvmField
-        val ALLOW_ROTATION =
-            backedUpItem(RotationHelper.ALLOW_ROTATION_PREFERENCE_KEY, Boolean::class.java) {
-                RotationHelper.getAllowRotationDefaultValue(DisplayController.INSTANCE.get(it).info)
-            }
-
-        // Preferences for widget configurations
-        @JvmField
-        val RECONFIGURABLE_WIDGET_EDUCATION_TIP_SEEN =
-            backedUpItem("launcher.reconfigurable_widget_education_tip_seen", false)
-
-        @JvmStatic
-        fun <T> backedUpItem(
-            sharedPrefKey: String,
-            defaultValue: T,
-            encryptionType: EncryptionType = EncryptionType.ENCRYPTED
-        ): ConstantItem<T> =
-            ConstantItem(sharedPrefKey, isBackedUp = true, defaultValue, encryptionType)
-
-        @JvmStatic
-        fun <T> backedUpItem(
-            sharedPrefKey: String,
-            type: Class<out T>,
-            encryptionType: EncryptionType = EncryptionType.ENCRYPTED,
-            defaultValueFromContext: (c: Context) -> T
-        ): ContextualItem<T> =
-            ContextualItem(
-                sharedPrefKey,
-                isBackedUp = true,
-                defaultValueFromContext,
-                encryptionType,
-                type
-            )
-
-        @JvmStatic
-        fun <T> nonRestorableItem(
-            sharedPrefKey: String,
-            defaultValue: T,
-            encryptionType: EncryptionType = EncryptionType.ENCRYPTED
-        ): ConstantItem<T> =
-            ConstantItem(sharedPrefKey, isBackedUp = false, defaultValue, encryptionType)
-
-        @Deprecated("Don't use shared preferences directly. Use other LauncherPref methods.")
-        @JvmStatic
-        fun getPrefs(context: Context): SharedPreferences {
-            // Use application context for shared preferences, so we use single cached instance
-            return context.applicationContext.getSharedPreferences(
-                SHARED_PREFERENCES_KEY,
-                MODE_PRIVATE
-            )
-        }
-
-        @Deprecated("Don't use shared preferences directly. Use other LauncherPref methods.")
-        @JvmStatic
-        fun getDevicePrefs(context: Context): SharedPreferences {
-            // Use application context for shared preferences, so we use a single cached instance
-            return context.applicationContext.getSharedPreferences(
-                DEVICE_PREFERENCES_KEY,
-                MODE_PRIVATE
-            )
-        }
-    }
 }
 
 abstract class Item {
@@ -395,7 +428,7 @@
     val defaultValue: T,
     override val encryptionType: EncryptionType,
     // The default value can be null. If so, the type needs to be explicitly stated, or else NPE
-    override val type: Class<out T> = defaultValue!!::class.java
+    override val type: Class<out T> = defaultValue!!::class.java,
 ) : Item() {
 
     fun get(c: Context): T = LauncherPrefs.get(c).get(this)
@@ -406,7 +439,7 @@
     override val isBackedUp: Boolean,
     private val defaultSupplier: (c: Context) -> T,
     override val encryptionType: EncryptionType,
-    override val type: Class<out T>
+    override val type: Class<out T>,
 ) : Item() {
     private var default: T? = null
 
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 87ac193..1d2d161 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -16,8 +16,12 @@
 
 package com.android.launcher3;
 
+import static android.util.Base64.NO_PADDING;
+import static android.util.Base64.NO_WRAP;
+
 import android.database.sqlite.SQLiteDatabase;
 import android.provider.BaseColumns;
+import android.util.Base64;
 
 import androidx.annotation.NonNull;
 
@@ -354,8 +358,17 @@
      * Launcher settings
      */
     public static final class Settings {
-        public static final String LAYOUT_DIGEST_KEY = "launcher3.layout.provider.blob";
+        public static final String LAYOUT_PROVIDER_KEY = "launcher3.layout.provider";
         public static final String LAYOUT_DIGEST_LABEL = "launcher-layout";
         public static final String LAYOUT_DIGEST_TAG = "ignore";
+        public static final String BLOB_KEY_PREFIX = "blob://";
+
+        /**
+         * Creates a key to be used for {@link #LAYOUT_PROVIDER_KEY}
+         * @param digest byte[] representing the message digest for the blob handle
+         */
+        public static String createBlobProviderKey(byte[] digest) {
+            return BLOB_KEY_PREFIX + Base64.encodeToString(digest, NO_WRAP | NO_PADDING);
+        }
     }
 }
diff --git a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
index 4e1e950..e705d94 100644
--- a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
@@ -346,17 +346,12 @@
                     (LetterListTextView) LayoutInflater.from(context).inflate(
                             R.layout.fast_scroller_letter_list_text_view, mLetterList, false);
             int viewId = View.generateViewId();
-            textView.setId(viewId);
+            textView.apply(sectionInfo /* FastScrollSectionInfo */, viewId /* viewId */);
             sectionInfo.setId(viewId);
-            textView.setText(sectionInfo.sectionName);
             if (i == fastScrollSections.size() - 1) {
                 // The last section info is just a duplicate so that user can scroll to the bottom.
                 textView.setVisibility(INVISIBLE);
             }
-            ConstraintLayout.LayoutParams lp = new ConstraintLayout.LayoutParams(
-                    MATCH_CONSTRAINT, WRAP_CONTENT);
-            lp.dimensionRatio = "v,1:1";
-            textView.setLayoutParams(lp);
             textViews.add(textView);
             mLetterList.addView(textView);
         }
diff --git a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
index 8e44d65..709b52a 100644
--- a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
+++ b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
@@ -106,6 +106,7 @@
     // The of ordered component names as a result of a search query
     private final ArrayList<AdapterItem> mSearchResults = new ArrayList<>();
     private final SpannableString mPrivateProfileAppScrollerBadge;
+    private final SpannableString mPrivateProfileDividerBadge;
     private BaseAllAppsAdapter<T> mAdapter;
     private AppInfoComparator mAppNameComparator;
     private int mNumAppsPerRowAllApps;
@@ -124,9 +125,14 @@
             mAllAppsStore.addUpdateListener(this);
         }
         mPrivateProfileAppScrollerBadge = new SpannableString(" ");
-        mPrivateProfileAppScrollerBadge.setSpan(new ImageSpan(context,
+        mPrivateProfileAppScrollerBadge.setSpan(new ImageSpan(context, Flags.letterFastScroller()
+                        ? R.drawable.ic_private_profile_letter_list_fast_scroller_badge :
                         R.drawable.ic_private_profile_app_scroller_badge, ImageSpan.ALIGN_CENTER),
                 0, 1, Spannable.SPAN_EXCLUSIVE_EXCLUSIVE);
+        mPrivateProfileDividerBadge = new SpannableString(" ");
+        mPrivateProfileDividerBadge.setSpan(new ImageSpan(context,
+                        R.drawable.ic_private_profile_divider_badge, ImageSpan.ALIGN_CENTER),
+                0, 1, Spannable.SPAN_EXCLUSIVE_EXCLUSIVE);
     }
 
     /** Set the number of apps per row when device profile changes. */
@@ -404,6 +410,11 @@
         // Add system apps separator.
         if (Flags.privateSpaceSysAppsSeparation()) {
             position = mPrivateProviderManager.addSystemAppsDivider(mAdapterItems);
+            if (Flags.letterFastScroller()) {
+                FastScrollSectionInfo sectionInfo =
+                        new FastScrollSectionInfo(mPrivateProfileDividerBadge, position);
+                mFastScrollerSections.add(sectionInfo);
+            }
         }
         // Add system apps.
         position = addAppsWithSections(split.get(false), position);
@@ -437,8 +448,11 @@
                 Log.d(TAG, "addAppsWithSections: adding sectionName: " + sectionName
                     + " with appInfoTitle: " + info.title);
                 lastSectionName = sectionName;
-                mFastScrollerSections.add(new FastScrollSectionInfo(hasPrivateApps ?
-                        mPrivateProfileAppScrollerBadge : sectionName, position));
+                boolean usePrivateAppScrollerBadge = !Flags.letterFastScroller() && hasPrivateApps;
+                FastScrollSectionInfo sectionInfo = new FastScrollSectionInfo(
+                        usePrivateAppScrollerBadge ?
+                                mPrivateProfileAppScrollerBadge : sectionName, position);
+                mFastScrollerSections.add(sectionInfo);
             }
             position++;
         }
diff --git a/src/com/android/launcher3/allapps/LetterListTextView.java b/src/com/android/launcher3/allapps/LetterListTextView.java
index 9326d79..433a7f2 100644
--- a/src/com/android/launcher3/allapps/LetterListTextView.java
+++ b/src/com/android/launcher3/allapps/LetterListTextView.java
@@ -16,6 +16,9 @@
 
 package com.android.launcher3.allapps;
 
+import static androidx.constraintlayout.widget.ConstraintSet.MATCH_CONSTRAINT;
+import static androidx.constraintlayout.widget.ConstraintSet.WRAP_CONTENT;
+
 import android.content.Context;
 import android.graphics.PorterDuff;
 import android.graphics.PorterDuffColorFilter;
@@ -23,6 +26,7 @@
 import android.util.AttributeSet;
 import android.widget.TextView;
 
+import androidx.constraintlayout.widget.ConstraintLayout;
 import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.R;
@@ -71,6 +75,20 @@
     }
 
     /**
+     * Applies a viewId to the letter list text view and sets the background and text based on the
+     * sectionInfo.
+     */
+    public void apply(AlphabeticalAppsList.FastScrollSectionInfo fastScrollSectionInfo,
+            int viewId) {
+        setId(viewId);
+        setText(fastScrollSectionInfo.sectionName);
+        ConstraintLayout.LayoutParams lp = new ConstraintLayout.LayoutParams(
+                MATCH_CONSTRAINT, WRAP_CONTENT);
+        lp.dimensionRatio = "v,1:1";
+        setLayoutParams(lp);
+    }
+
+    /**
      * Animates the letter list text view based on the current finger position.
      *
      * @param currentFingerY The Y position of where the finger is placed on the fastScroller in
diff --git a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
index d39c5de..088277b 100644
--- a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
+++ b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
@@ -21,6 +21,7 @@
 import com.android.launcher3.pm.InstallSessionHelper;
 import com.android.launcher3.util.DaggerSingletonTracker;
 import com.android.launcher3.util.ScreenOnTracker;
+import com.android.launcher3.util.SettingsCache;
 
 import dagger.BindsInstance;
 
@@ -36,6 +37,7 @@
     DaggerSingletonTracker getDaggerSingletonTracker();
     InstallSessionHelper getInstallSessionHelper();
     ScreenOnTracker getScreenOnTracker();
+    SettingsCache getSettingsCache();
 
     /** Builder for LauncherBaseAppComponent. */
     interface Builder {
diff --git a/src/com/android/launcher3/dragndrop/AddItemActivity.java b/src/com/android/launcher3/dragndrop/AddItemActivity.java
index a3cfe5c..25de479 100644
--- a/src/com/android/launcher3/dragndrop/AddItemActivity.java
+++ b/src/com/android/launcher3/dragndrop/AddItemActivity.java
@@ -281,7 +281,7 @@
                 new PinShortcutRequestActivityInfo(mRequest, this);
         mWidgetCell.getWidgetView().setTag(new PendingAddShortcutInfo(shortcutInfo));
         applyWidgetItemAsync(
-                () -> new WidgetItem(shortcutInfo, mApp.getIconCache(), getPackageManager()));
+                () -> new WidgetItem(shortcutInfo, mApp.getIconCache()));
         return new PackageItemInfo(mRequest.getShortcutInfo().getPackage(),
                 mRequest.getShortcutInfo().getUserHandle());
     }
diff --git a/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java b/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java
index cc5e890..a6a50d7 100644
--- a/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java
+++ b/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java
@@ -30,7 +30,6 @@
 import android.content.Intent;
 import android.content.pm.LauncherApps;
 import android.content.pm.LauncherApps.PinItemRequest;
-import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
 import android.graphics.drawable.Drawable;
 import android.os.Build;
@@ -40,7 +39,7 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
-import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.BaseIconCache;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.pm.PinRequestHelper;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
@@ -82,12 +81,12 @@
     }
 
     @Override
-    public CharSequence getLabel(PackageManager pm) {
+    public CharSequence getLabel() {
         return mInfo.getShortLabel();
     }
 
     @Override
-    public Drawable getFullResIcon(IconCache cache) {
+    public Drawable getFullResIcon(BaseIconCache cache) {
         Drawable d = mContext.getSystemService(LauncherApps.class)
                 .getShortcutIconDrawable(mInfo, LauncherAppState.getIDP(mContext).fillResIconDpi);
         if (d == null) {
diff --git a/src/com/android/launcher3/folder/Folder.java b/src/com/android/launcher3/folder/Folder.java
index 7bec768..5defef3 100644
--- a/src/com/android/launcher3/folder/Folder.java
+++ b/src/com/android/launcher3/folder/Folder.java
@@ -317,10 +317,7 @@
                 | InputType.TYPE_TEXT_FLAG_NO_SUGGESTIONS
                 | InputType.TYPE_TEXT_FLAG_CAP_WORDS);
         mFolderName.forceDisableSuggestions(true);
-        mFolderName.setPadding(mFolderName.getPaddingLeft(),
-                (getFooterHeight() - mFolderName.getLineHeight()) / 2,
-                mFolderName.getPaddingRight(),
-                (getFooterHeight() - mFolderName.getLineHeight()) / 2);
+
 
         mKeyboardInsetAnimationCallback = new KeyboardInsetAnimationCallback(this);
         setWindowInsetsAnimationCallback(mKeyboardInsetAnimationCallback);
diff --git a/src/com/android/launcher3/folder/FolderPagedView.java b/src/com/android/launcher3/folder/FolderPagedView.java
index 9dc2d24..fe26194 100644
--- a/src/com/android/launcher3/folder/FolderPagedView.java
+++ b/src/com/android/launcher3/folder/FolderPagedView.java
@@ -373,8 +373,9 @@
         // Update footer
         mPageIndicator.setVisibility(getPageCount() > 1 ? View.VISIBLE : View.GONE);
         // Set the gravity as LEFT or RIGHT instead of START, as START depends on the actual text.
-        mFolder.getFolderName().setGravity(getPageCount() > 1
-                ? (mIsRtl ? Gravity.RIGHT : Gravity.LEFT) : Gravity.CENTER_HORIZONTAL);
+        int horizontalGravity = getPageCount() > 1
+                ? (mIsRtl ? Gravity.RIGHT : Gravity.LEFT) : Gravity.CENTER_HORIZONTAL;
+        mFolder.getFolderName().setGravity(horizontalGravity | Gravity.CENTER_VERTICAL);
     }
 
     public int getDesiredWidth() {
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index b6fe66a..e7c4024 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -51,8 +51,8 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.icons.cache.BaseIconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.icons.cache.CachedObjectCachingLogic;
-import com.android.launcher3.icons.cache.CachingLogic;
 import com.android.launcher3.icons.cache.LauncherActivityCachingLogic;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.AppInfo;
@@ -92,9 +92,6 @@
     private final Predicate<ItemInfoWithIcon> mIsUsingFallbackOrNonDefaultIconCheck = w ->
             w.bitmap != null && (w.bitmap.isNullOrLowRes() || !isDefaultIcon(w.bitmap, w.user));
 
-    private final CachingLogic<ComponentWithLabel> mComponentWithLabelCachingLogic;
-    private final CachingLogic<LauncherActivityInfo> mLauncherActivityInfoCachingLogic;
-
     private final LauncherApps mLauncherApps;
     private final UserCache mUserManager;
     private final InstantAppResolver mInstantAppResolver;
@@ -108,8 +105,6 @@
             IconProvider iconProvider) {
         super(context, dbFileName, MODEL_EXECUTOR.getLooper(),
                 idp.fillResIconDpi, idp.iconBitmapSize, true /* inMemoryCache */, iconProvider);
-        mComponentWithLabelCachingLogic = new CachedObjectCachingLogic(context);
-        mLauncherActivityInfoCachingLogic = LauncherActivityCachingLogic.INSTANCE;
         mLauncherApps = mContext.getSystemService(LauncherApps.class);
         mUserManager = UserCache.INSTANCE.get(mContext);
         mInstantAppResolver = InstantAppResolver.newInstance(mContext);
@@ -143,7 +138,7 @@
         removeIconsForPkg(packageName, user);
         long userSerial = mUserManager.getSerialNumberForUser(user);
         for (LauncherActivityInfo app : mLauncherApps.getActivityList(packageName, user)) {
-            addIconToDBAndMemCache(app, mLauncherActivityInfoCachingLogic, userSerial);
+            addIconToDBAndMemCache(app, LauncherActivityCachingLogic.INSTANCE, userSerial);
         }
     }
 
@@ -211,7 +206,7 @@
      */
     public synchronized void updateTitleAndIcon(AppInfo application) {
         CacheEntry entry = cacheLocked(application.componentName,
-                application.user, () -> null, mLauncherActivityInfoCachingLogic,
+                application.user, () -> null, LauncherActivityCachingLogic.INSTANCE,
                 application.usingLowResIcon() ? LookupFlag.USE_LOW_RES : LookupFlag.DEFAULT);
         if (entry.bitmap != null || !isDefaultIcon(entry.bitmap, application.user)) {
             applyCacheEntry(entry, application);
@@ -326,9 +321,9 @@
     /**
      * Loads and returns the icon for the provided object without adding it to memCache
      */
-    public synchronized String getTitleNoCache(ComponentWithLabel info) {
+    public synchronized String getTitleNoCache(CachedObject info) {
         CacheEntry entry = cacheLocked(info.getComponent(), info.getUser(), () -> info,
-                mComponentWithLabelCachingLogic,
+                CachedObjectCachingLogic.INSTANCE,
                 LookupFlag.USE_LOW_RES | LookupFlag.SKIP_ADD_TO_MEM_CACHE);
         return Utilities.trim(entry.title);
     }
@@ -344,7 +339,7 @@
         if (usePkgIcon) lookupFlags |= LookupFlag.USE_PACKAGE_ICON;
         if (useLowResIcon) lookupFlags |= LookupFlag.USE_LOW_RES;
         CacheEntry entry = cacheLocked(infoInOut.getTargetComponent(), infoInOut.user,
-                activityInfoProvider, mLauncherActivityInfoCachingLogic, lookupFlags);
+                activityInfoProvider, LauncherActivityCachingLogic.INSTANCE, lookupFlags);
         applyCacheEntry(entry, infoInOut);
     }
 
@@ -445,7 +440,7 @@
                                 cn,
                                 /* user = */ sectionKey.first,
                                 () -> duplicateIconRequests.get(0).launcherActivityInfo,
-                                mLauncherActivityInfoCachingLogic,
+                                LauncherActivityCachingLogic.INSTANCE,
                                 sectionKey.second ? LookupFlag.USE_LOW_RES : LookupFlag.DEFAULT,
                                 c);
 
@@ -494,7 +489,7 @@
                     loadFallbackIcon(
                             lai,
                             entry,
-                            mLauncherActivityInfoCachingLogic,
+                            LauncherActivityCachingLogic.INSTANCE,
                             /* usePackageIcon= */ false,
                             /* usePackageTitle= */ loadFallbackTitle,
                             cn,
@@ -504,7 +499,7 @@
                     loadFallbackTitle(
                             lai,
                             entry,
-                            mLauncherActivityInfoCachingLogic,
+                            LauncherActivityCachingLogic.INSTANCE,
                             sectionKey.first);
                 }
 
diff --git a/src/com/android/launcher3/icons/Legacy.kt b/src/com/android/launcher3/icons/Legacy.kt
deleted file mode 100644
index 3bf3bb2..0000000
--- a/src/com/android/launcher3/icons/Legacy.kt
+++ /dev/null
@@ -1,31 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.icons
-
-import com.android.launcher3.icons.cache.CachedObject
-
-/**
- * This files contains some definitions used during refactoring to avoid breaking changes.
- *
- * TODO(b/366237794) remove this file once refactoring is complete
- */
-
-/** Temporary interface to allow easier refactoring */
-interface ComponentWithLabel : CachedObject<IconCache>
-
-/** Temporary interface to allow easier refactoring */
-interface ComponentWithLabelAndIcon : ComponentWithLabel
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index dff5463..09d1146 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -72,8 +72,8 @@
 import com.android.launcher3.folder.FolderNameProvider;
 import com.android.launcher3.icons.CacheableShortcutCachingLogic;
 import com.android.launcher3.icons.CacheableShortcutInfo;
-import com.android.launcher3.icons.ComponentWithLabelAndIcon;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.icons.cache.CachedObjectCachingLogic;
 import com.android.launcher3.icons.cache.IconCacheUpdateHandler;
 import com.android.launcher3.icons.cache.LauncherActivityCachingLogic;
@@ -335,8 +335,7 @@
             verifyNotStopped();
 
             // fourth step
-            List<ComponentWithLabelAndIcon> allWidgetsList =
-                    mBgDataModel.widgetsModel.update(mApp, null);
+            List<CachedObject> allWidgetsList = mBgDataModel.widgetsModel.update(mApp, null);
             logASplit("load widgets");
 
             verifyNotStopped();
@@ -364,7 +363,7 @@
             }
 
             updateHandler.updateIcons(allWidgetsList,
-                    new CachedObjectCachingLogic(mApp.getContext()),
+                    CachedObjectCachingLogic.INSTANCE,
                     mApp.getModel()::onWidgetLabelsUpdated);
             logASplit("save widgets in icon cache");
 
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index da1a221..5d66d16 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -25,7 +25,8 @@
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.addTableToDb;
-import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_KEY;
+import static com.android.launcher3.LauncherSettings.Settings.BLOB_KEY_PREFIX;
+import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_PROVIDER_KEY;
 import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_LABEL;
 import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_TAG;
 import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
@@ -548,9 +549,15 @@
     private AutoInstallsLayout createWorkspaceLoaderFromAppRestriction(
             LauncherWidgetHolder widgetHolder) {
         ContentResolver cr = mContext.getContentResolver();
-        String blobHandlerDigest = Settings.Secure.getString(cr, LAYOUT_DIGEST_KEY);
-        if (!TextUtils.isEmpty(blobHandlerDigest)) {
+        String systemLayoutProvider = Settings.Secure.getString(cr, LAYOUT_PROVIDER_KEY);
+        if (TextUtils.isEmpty(systemLayoutProvider)) {
+            return null;
+        }
+
+        // Try the blob store first
+        if (systemLayoutProvider.startsWith(BLOB_KEY_PREFIX)) {
             BlobStoreManager blobManager = mContext.getSystemService(BlobStoreManager.class);
+            String blobHandlerDigest = systemLayoutProvider.substring(BLOB_KEY_PREFIX.length());
             try (InputStream in = new ParcelFileDescriptor.AutoCloseInputStream(
                     blobManager.openBlob(BlobHandle.createWithSha256(
                             Base64.decode(blobHandlerDigest, NO_WRAP | NO_PADDING),
@@ -562,25 +569,21 @@
             }
         }
 
-        String authority = Settings.Secure.getString(cr, "launcher3.layout.provider");
-        if (TextUtils.isEmpty(authority)) {
-            return null;
-        }
-
+        // Try contentProvider based provider
         PackageManager pm = mContext.getPackageManager();
-        ProviderInfo pi = pm.resolveContentProvider(authority, 0);
+        ProviderInfo pi = pm.resolveContentProvider(systemLayoutProvider, 0);
         if (pi == null) {
-            Log.e(TAG, "No provider found for authority " + authority);
+            Log.e(TAG, "No provider found for authority " + systemLayoutProvider);
             return null;
         }
-        Uri uri = getLayoutUri(authority, mContext);
+        Uri uri = getLayoutUri(systemLayoutProvider, mContext);
         try (InputStream in = cr.openInputStream(uri)) {
-            Log.d(TAG, "Loading layout from " + authority);
+            Log.d(TAG, "Loading layout from " + systemLayoutProvider);
 
             Resources res = pm.getResourcesForApplication(pi.applicationInfo);
             return getAutoInstallsLayoutFromIS(in, widgetHolder, SourceResources.wrap(res));
         } catch (Exception e) {
-            Log.e(TAG, "Error getting layout stream from: " + authority , e);
+            Log.e(TAG, "Error getting layout stream from: " + systemLayoutProvider , e);
             return null;
         }
     }
diff --git a/src/com/android/launcher3/model/ModelLauncherCallbacks.kt b/src/com/android/launcher3/model/ModelLauncherCallbacks.kt
index 2ee5b80..7ba2dad 100644
--- a/src/com/android/launcher3/model/ModelLauncherCallbacks.kt
+++ b/src/com/android/launcher3/model/ModelLauncherCallbacks.kt
@@ -17,10 +17,12 @@
 package com.android.launcher3.model
 
 import android.content.pm.LauncherApps
+import android.content.pm.PackageInstaller.SessionInfo
 import android.content.pm.ShortcutInfo
 import android.os.UserHandle
 import android.text.TextUtils
 import com.android.launcher3.LauncherModel.ModelUpdateTask
+import com.android.launcher3.config.FeatureFlags
 import com.android.launcher3.logging.FileLog
 import com.android.launcher3.model.PackageUpdatedTask.OP_ADD
 import com.android.launcher3.model.PackageUpdatedTask.OP_REMOVE
@@ -28,6 +30,9 @@
 import com.android.launcher3.model.PackageUpdatedTask.OP_UNAVAILABLE
 import com.android.launcher3.model.PackageUpdatedTask.OP_UNSUSPEND
 import com.android.launcher3.model.PackageUpdatedTask.OP_UPDATE
+import com.android.launcher3.pm.InstallSessionTracker
+import com.android.launcher3.pm.PackageInstallInfo
+import com.android.launcher3.util.PackageUserKey
 import java.util.function.Consumer
 
 /**
@@ -35,7 +40,7 @@
  * model tasks
  */
 class ModelLauncherCallbacks(private var taskExecutor: Consumer<ModelUpdateTask>) :
-    LauncherApps.Callback() {
+    LauncherApps.Callback(), InstallSessionTracker.Callback {
 
     override fun onPackageAdded(packageName: String, user: UserHandle) {
         FileLog.d(TAG, "onPackageAdded triggered for packageName=$packageName, user=$user")
@@ -49,7 +54,7 @@
     override fun onPackageLoadingProgressChanged(
         packageName: String,
         user: UserHandle,
-        progress: Float
+        progress: Float,
     ) {
         taskExecutor.accept(PackageIncrementalDownloadUpdatedTask(packageName, user, progress))
     }
@@ -62,7 +67,7 @@
     override fun onPackagesAvailable(
         vararg packageNames: String,
         user: UserHandle,
-        replacing: Boolean
+        replacing: Boolean,
     ) {
         taskExecutor.accept(PackageUpdatedTask(OP_UPDATE, user, *packageNames))
     }
@@ -74,7 +79,7 @@
     override fun onPackagesUnavailable(
         packageNames: Array<String>,
         user: UserHandle,
-        replacing: Boolean
+        replacing: Boolean,
     ) {
         if (!replacing) {
             taskExecutor.accept(PackageUpdatedTask(OP_UNAVAILABLE, user, *packageNames))
@@ -88,7 +93,7 @@
     override fun onShortcutsChanged(
         packageName: String,
         shortcuts: MutableList<ShortcutInfo>,
-        user: UserHandle
+        user: UserHandle,
     ) {
         taskExecutor.accept(ShortcutsChangedTask(packageName, shortcuts, user, true))
     }
@@ -98,6 +103,37 @@
         taskExecutor.accept(PackageUpdatedTask(OP_REMOVE, user, *packages.toTypedArray()))
     }
 
+    override fun onSessionFailure(packageName: String, user: UserHandle) {
+        taskExecutor.accept(SessionFailureTask(packageName, user))
+    }
+
+    override fun onPackageStateChanged(installInfo: PackageInstallInfo) {
+        taskExecutor.accept(PackageInstallStateChangedTask(installInfo))
+    }
+
+    override fun onUpdateSessionDisplay(key: PackageUserKey, info: SessionInfo) {
+        /** Updates the icons and label of all pending icons for the provided package name. */
+        taskExecutor.accept { controller, _, _ ->
+            controller.app.iconCache.updateSessionCache(key, info)
+        }
+        taskExecutor.accept(
+            CacheDataUpdatedTask(
+                CacheDataUpdatedTask.OP_SESSION_UPDATE,
+                key.mUser,
+                hashSetOf(key.mPackageName),
+            )
+        )
+    }
+
+    override fun onInstallSessionCreated(sessionInfo: PackageInstallInfo) {
+        if (FeatureFlags.PROMISE_APPS_IN_ALL_APPS.get()) {
+            taskExecutor.accept { taskController, _, apps ->
+                apps.addPromiseApp(taskController.app.context, sessionInfo)
+                taskController.bindApplicationsIfNeeded()
+            }
+        }
+    }
+
     companion object {
         private const val TAG = "LauncherAppsCallbackImpl"
     }
diff --git a/src/com/android/launcher3/model/SessionFailureTask.kt b/src/com/android/launcher3/model/SessionFailureTask.kt
new file mode 100644
index 0000000..0d006fa
--- /dev/null
+++ b/src/com/android/launcher3/model/SessionFailureTask.kt
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.model
+
+import android.content.ComponentName
+import android.os.UserHandle
+import android.text.TextUtils
+import com.android.launcher3.LauncherModel.ModelUpdateTask
+import com.android.launcher3.icons.cache.BaseIconCache
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.util.ApplicationInfoWrapper
+import com.android.launcher3.util.ItemInfoMatcher
+
+/** Model task run when there is a package install session failure */
+class SessionFailureTask(val packageName: String, val user: UserHandle) : ModelUpdateTask {
+
+    override fun execute(
+        taskController: ModelTaskController,
+        dataModel: BgDataModel,
+        apps: AllAppsList,
+    ) {
+        val iconCache = taskController.app.iconCache
+        val isAppArchived =
+            ApplicationInfoWrapper(taskController.app.context, packageName, user).isArchived()
+        synchronized(dataModel) {
+            if (isAppArchived) {
+                val updatedItems = mutableListOf<WorkspaceItemInfo>()
+                // Remove package icon cache entry for archived app in case of a session
+                // failure.
+                iconCache.remove(
+                    ComponentName(packageName, packageName + BaseIconCache.EMPTY_CLASS_NAME),
+                    user,
+                )
+                for (info in dataModel.itemsIdMap) {
+                    if (info is WorkspaceItemInfo && info.isArchived && user == info.user) {
+                        // Refresh icons on the workspace for archived apps.
+                        iconCache.getTitleAndIcon(info, info.usingLowResIcon())
+                        updatedItems.add(info)
+                    }
+                }
+
+                if (updatedItems.isNotEmpty()) {
+                    taskController.bindUpdatedWorkspaceItems(updatedItems)
+                }
+                apps.updateIconsAndLabels(hashSetOf(packageName), user)
+                taskController.bindApplicationsIfNeeded()
+            } else {
+                val removedItems =
+                    dataModel.itemsIdMap.filter { info ->
+                        (info is WorkspaceItemInfo && info.hasPromiseIconUi()) &&
+                            user == info.user &&
+                            TextUtils.equals(packageName, info.intent.getPackage())
+                    }
+                if (removedItems.isNotEmpty()) {
+                    taskController.deleteAndBindComponentsRemoved(
+                        ItemInfoMatcher.ofItems(removedItems),
+                        "removed because install session failed",
+                    )
+                }
+            }
+        }
+    }
+}
diff --git a/src/com/android/launcher3/model/WidgetItem.java b/src/com/android/launcher3/model/WidgetItem.java
index ac9f2d6..e757a68 100644
--- a/src/com/android/launcher3/model/WidgetItem.java
+++ b/src/com/android/launcher3/model/WidgetItem.java
@@ -7,7 +7,6 @@
 import android.annotation.SuppressLint;
 import android.content.Context;
 import android.content.pm.ActivityInfo;
-import android.content.pm.PackageManager;
 import android.content.res.Resources;
 import android.util.SparseArray;
 import android.widget.RemoteViews;
@@ -75,10 +74,10 @@
         this(info, idp, iconCache, context, new WidgetManagerHelper(context));
     }
 
-    public WidgetItem(ShortcutConfigActivityInfo info, IconCache iconCache, PackageManager pm) {
+    public WidgetItem(ShortcutConfigActivityInfo info, IconCache iconCache) {
         super(info.getComponent(), info.getUser());
         label = info.isPersistable() ? iconCache.getTitleNoCache(info) :
-                Utilities.trim(info.getLabel(pm));
+                Utilities.trim(info.getLabel());
         description = null;
         widgetInfo = null;
         activityInfo = info;
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index c949ce6..b450f46 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -14,7 +14,6 @@
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.os.UserHandle;
 import android.util.Log;
 import android.util.Pair;
@@ -27,8 +26,8 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.icons.ComponentWithLabelAndIcon;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
 import com.android.launcher3.util.ComponentKey;
@@ -96,20 +95,18 @@
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
      */
-    public List<ComponentWithLabelAndIcon> update(
+    public List<CachedObject> update(
             LauncherAppState app, @Nullable PackageUserKey packageUser) {
         if (!WIDGETS_ENABLED) {
-            return Collections.emptyList();
+            return new ArrayList<>();
         }
         Preconditions.assertWorkerThread();
 
         Context context = app.getContext();
         final ArrayList<WidgetItem> widgetsAndShortcuts = new ArrayList<>();
-        List<ComponentWithLabelAndIcon> updatedItems = new ArrayList<>();
+        List<CachedObject> updatedItems = new ArrayList<>();
         try {
             InvariantDeviceProfile idp = app.getInvariantDeviceProfile();
-            PackageManager pm = app.getContext().getPackageManager();
-
             // Widgets
             WidgetManagerHelper widgetManager = new WidgetManagerHelper(context);
             for (AppWidgetProviderInfo widgetInfo : widgetManager.getAllProviders(packageUser)) {
@@ -125,7 +122,7 @@
             // Shortcuts
             for (ShortcutConfigActivityInfo info :
                     queryList(context, packageUser)) {
-                widgetsAndShortcuts.add(new WidgetItem(info, app.getIconCache(), pm));
+                widgetsAndShortcuts.add(new WidgetItem(info, app.getIconCache()));
                 updatedItems.add(info);
             }
             setWidgetsAndShortcuts(widgetsAndShortcuts, app, packageUser);
@@ -190,8 +187,7 @@
                     WidgetItem item = items.get(i);
                     if (item.user.equals(user)) {
                         if (item.activityInfo != null) {
-                            items.set(i, new WidgetItem(item.activityInfo, app.getIconCache(),
-                                    app.getContext().getPackageManager()));
+                            items.set(i, new WidgetItem(item.activityInfo, app.getIconCache()));
                         } else {
                             items.set(i, new WidgetItem(item.widgetInfo,
                                     app.getInvariantDeviceProfile(), app.getIconCache(),
diff --git a/src/com/android/launcher3/model/data/AppPairInfo.kt b/src/com/android/launcher3/model/data/AppPairInfo.kt
index 2eb6154..3496c17 100644
--- a/src/com/android/launcher3/model/data/AppPairInfo.kt
+++ b/src/com/android/launcher3/model/data/AppPairInfo.kt
@@ -74,7 +74,7 @@
             (ActivityContext.lookupContext(context) as ActivityContext).getDeviceProfile().isTablet
         return Pair(
             isTablet || !getFirstApp().isNonResizeable(),
-            isTablet || !getSecondApp().isNonResizeable()
+            isTablet || !getSecondApp().isNonResizeable(),
         )
     }
 
@@ -105,10 +105,10 @@
     }
 
     /** Generates an ItemInfo for logging. */
-    override fun buildProto(cInfo: CollectionInfo?): LauncherAtom.ItemInfo {
+    override fun buildProto(cInfo: CollectionInfo?, context: Context): LauncherAtom.ItemInfo {
         val appPairIcon = LauncherAtom.FolderIcon.newBuilder().setCardinality(contents.size)
         appPairIcon.setLabelInfo(title.toString())
-        return getDefaultItemInfoBuilder()
+        return getDefaultItemInfoBuilder(context)
             .setFolderIcon(appPairIcon)
             .setRank(rank)
             .setContainerInfo(getContainerInfo())
diff --git a/src/com/android/launcher3/model/data/CollectionInfo.kt b/src/com/android/launcher3/model/data/CollectionInfo.kt
index 4f5e12f..12ba164 100644
--- a/src/com/android/launcher3/model/data/CollectionInfo.kt
+++ b/src/com/android/launcher3/model/data/CollectionInfo.kt
@@ -17,7 +17,6 @@
 package com.android.launcher3.model.data
 
 import com.android.launcher3.LauncherSettings
-import com.android.launcher3.logger.LauncherAtom
 import com.android.launcher3.util.ContentWriter
 import java.util.function.Predicate
 
@@ -42,9 +41,4 @@
         super.onAddToDatabase(writer)
         writer.put(LauncherSettings.Favorites.TITLE, title)
     }
-
-    /** Returns the collection wrapped as {@link LauncherAtom.ItemInfo} for logging. */
-    override fun buildProto(): LauncherAtom.ItemInfo {
-        return buildProto(null)
-    }
 }
diff --git a/src/com/android/launcher3/model/data/FolderInfo.java b/src/com/android/launcher3/model/data/FolderInfo.java
index 18d2b85..f0f2892 100644
--- a/src/com/android/launcher3/model/data/FolderInfo.java
+++ b/src/com/android/launcher3/model/data/FolderInfo.java
@@ -24,6 +24,8 @@
 import static com.android.launcher3.logger.LauncherAtom.Attribute.MANUAL_LABEL;
 import static com.android.launcher3.logger.LauncherAtom.Attribute.SUGGESTED_LABEL;
 
+import android.content.Context;
+
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
@@ -245,13 +247,13 @@
 
     @NonNull
     @Override
-    public LauncherAtom.ItemInfo buildProto(@Nullable CollectionInfo cInfo) {
+    public LauncherAtom.ItemInfo buildProto(@Nullable CollectionInfo cInfo, Context context) {
         FolderIcon.Builder folderIcon = FolderIcon.newBuilder()
                 .setCardinality(getContents().size());
         if (LabelState.SUGGESTED.equals(getLabelState())) {
             folderIcon.setLabelInfo(title.toString());
         }
-        return getDefaultItemInfoBuilder()
+        return getDefaultItemInfoBuilder(context)
                 .setFolderIcon(folderIcon)
                 .setRank(rank)
                 .addItemAttributes(getLabelState().mLogAttribute)
diff --git a/src/com/android/launcher3/model/data/ItemInfo.java b/src/com/android/launcher3/model/data/ItemInfo.java
index b706d24..c22a8a5 100644
--- a/src/com/android/launcher3/model/data/ItemInfo.java
+++ b/src/com/android/launcher3/model/data/ItemInfo.java
@@ -36,6 +36,7 @@
 
 import android.content.ComponentName;
 import android.content.ContentValues;
+import android.content.Context;
 import android.content.Intent;
 import android.net.Uri;
 import android.os.Process;
@@ -352,16 +353,16 @@
      * Creates {@link LauncherAtom.ItemInfo} with important fields and parent container info.
      */
     @NonNull
-    public LauncherAtom.ItemInfo buildProto() {
-        return buildProto(null);
+    public LauncherAtom.ItemInfo buildProto(Context context) {
+        return buildProto(null, context);
     }
 
     /**
      * Creates {@link LauncherAtom.ItemInfo} with important fields and parent container info.
      */
     @NonNull
-    public LauncherAtom.ItemInfo buildProto(@Nullable final CollectionInfo cInfo) {
-        LauncherAtom.ItemInfo.Builder itemBuilder = getDefaultItemInfoBuilder();
+    public LauncherAtom.ItemInfo buildProto(@Nullable final CollectionInfo cInfo, Context context) {
+        LauncherAtom.ItemInfo.Builder itemBuilder = getDefaultItemInfoBuilder(context);
         Optional<ComponentName> nullableComponent = Optional.ofNullable(getTargetComponent());
         switch (itemType) {
             case ITEM_TYPE_APPLICATION:
@@ -434,10 +435,10 @@
     }
 
     @NonNull
-    protected LauncherAtom.ItemInfo.Builder getDefaultItemInfoBuilder() {
+    protected LauncherAtom.ItemInfo.Builder getDefaultItemInfoBuilder(Context context) {
         LauncherAtom.ItemInfo.Builder itemBuilder = LauncherAtom.ItemInfo.newBuilder();
-        SettingsCache.INSTANCE.executeIfCreated(cache ->
-                itemBuilder.setIsKidsMode(cache.getValue(NAV_BAR_KIDS_MODE, 0)));
+        itemBuilder.setIsKidsMode(
+                SettingsCache.INSTANCE.get(context).getValue(NAV_BAR_KIDS_MODE, 0));
         UserCache.INSTANCE.executeIfCreated(cache ->
                 itemBuilder.setUserType(getUserType(cache.getUserInfo(user))));
         itemBuilder.setRank(rank);
diff --git a/src/com/android/launcher3/model/data/LauncherAppWidgetInfo.java b/src/com/android/launcher3/model/data/LauncherAppWidgetInfo.java
index 361f09d..7569ed5 100644
--- a/src/com/android/launcher3/model/data/LauncherAppWidgetInfo.java
+++ b/src/com/android/launcher3/model/data/LauncherAppWidgetInfo.java
@@ -24,6 +24,7 @@
 
 import android.appwidget.AppWidgetHostView;
 import android.content.ComponentName;
+import android.content.Context;
 import android.content.Intent;
 import android.content.res.Resources;
 import android.os.Process;
@@ -270,8 +271,9 @@
 
     @NonNull
     @Override
-    public LauncherAtom.ItemInfo buildProto(@Nullable CollectionInfo collectionInfo) {
-        LauncherAtom.ItemInfo info = super.buildProto(collectionInfo);
+    public LauncherAtom.ItemInfo buildProto(
+            @Nullable CollectionInfo collectionInfo, Context context) {
+        LauncherAtom.ItemInfo info = super.buildProto(collectionInfo, context);
         return info.toBuilder()
                 .setWidget(info.getWidget().toBuilder().setWidgetFeatures(widgetFeatures))
                 .addItemAttributes(getAttribute(sourceContainer))
diff --git a/src/com/android/launcher3/pm/ShortcutConfigActivityInfo.java b/src/com/android/launcher3/pm/ShortcutConfigActivityInfo.java
index 3064abf..409174e 100644
--- a/src/com/android/launcher3/pm/ShortcutConfigActivityInfo.java
+++ b/src/com/android/launcher3/pm/ShortcutConfigActivityInfo.java
@@ -29,7 +29,6 @@
 import android.content.pm.ApplicationInfo;
 import android.content.pm.LauncherActivityInfo;
 import android.content.pm.LauncherApps;
-import android.content.pm.PackageManager;
 import android.graphics.drawable.Drawable;
 import android.os.Build;
 import android.os.Process;
@@ -41,8 +40,8 @@
 
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
-import com.android.launcher3.icons.ComponentWithLabelAndIcon;
-import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.BaseIconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.ApplicationInfoWrapper;
 import com.android.launcher3.util.PackageUserKey;
@@ -54,7 +53,7 @@
 /**
  * Wrapper class for representing a shortcut configure activity.
  */
-public abstract class ShortcutConfigActivityInfo implements ComponentWithLabelAndIcon {
+public abstract class ShortcutConfigActivityInfo implements CachedObject {
 
     private static final String TAG = "SCActivityInfo";
 
@@ -91,7 +90,7 @@
     }
 
     @Override
-    public abstract Drawable getFullResIcon(IconCache cache);
+    public abstract Drawable getFullResIcon(BaseIconCache cache);
 
     /**
      * Return a WorkspaceItemInfo, if it can be created directly on drop, without requiring any
@@ -125,7 +124,7 @@
     }
 
     /**
-     * Returns true if various properties ({@link #getLabel(PackageManager)},
+     * Returns true if various properties ({@link #getLabel()},
      * {@link #getFullResIcon}) can be safely persisted.
      */
     public boolean isPersistable() {
@@ -144,12 +143,12 @@
         }
 
         @Override
-        public CharSequence getLabel(PackageManager pm) {
+        public CharSequence getLabel() {
             return mInfo.getLabel();
         }
 
         @Override
-        public Drawable getFullResIcon(IconCache cache) {
+        public Drawable getFullResIcon(BaseIconCache cache) {
             return cache.getFullResIcon(mInfo.getActivityInfo());
         }
 
diff --git a/src/com/android/launcher3/statemanager/StateManager.java b/src/com/android/launcher3/statemanager/StateManager.java
index 303290d..763f3ba 100644
--- a/src/com/android/launcher3/statemanager/StateManager.java
+++ b/src/com/android/launcher3/statemanager/StateManager.java
@@ -18,6 +18,7 @@
 
 import static android.animation.ValueAnimator.areAnimatorsEnabled;
 
+import static com.android.launcher3.Flags.enableStateManagerProtoLog;
 import static com.android.launcher3.anim.AnimatorPlaybackController.callListenerCommandRecursively;
 import static com.android.launcher3.states.StateAnimationConfig.HANDLE_STATE_APPLY;
 import static com.android.launcher3.states.StateAnimationConfig.SKIP_ALL_ANIMATIONS;
@@ -39,6 +40,7 @@
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.states.StateAnimationConfig.AnimationFlags;
 import com.android.launcher3.states.StateAnimationConfig.AnimationPropertyFlags;
+import com.android.launcher3.util.StateManagerProtoLogProxy;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -243,7 +245,10 @@
 
     private void goToState(
             STATE_TYPE state, boolean animated, long delay, AnimatorListener listener) {
-        if (DEBUG) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logGoToState(
+                    mState, state, getTrimmedStackTrace("StateManager.goToState"));
+        } else if (DEBUG) {
             Log.d(TAG, "goToState - fromState: " + mState + ", toState: " + state
                     + ", partial trace:\n" + getTrimmedStackTrace("StateManager.goToState"));
         }
@@ -331,7 +336,10 @@
      */
     public AnimatorSet createAtomicAnimation(
             STATE_TYPE fromState, STATE_TYPE toState, StateAnimationConfig config) {
-        if (DEBUG) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logCreateAtomicAnimation(
+                    mState, toState, getTrimmedStackTrace("StateManager.createAtomicAnimation"));
+        } else if (DEBUG) {
             Log.d(TAG, "createAtomicAnimation - fromState: " + fromState + ", toState: " + toState
                     + ", partial trace:\n" + getTrimmedStackTrace(
                             "StateManager.createAtomicAnimation"));
@@ -408,7 +416,9 @@
         mState = state;
         mStatefulContainer.onStateSetStart(mState);
 
-        if (DEBUG) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logOnStateTransitionStart(state);
+        } else if (DEBUG) {
             Log.d(TAG, "onStateTransitionStart - state: " + state);
         }
         for (int i = mListeners.size() - 1; i >= 0; i--) {
@@ -428,7 +438,9 @@
             setRestState(null);
         }
 
-        if (DEBUG) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logOnStateTransitionEnd(state);
+        } else if (DEBUG) {
             Log.d(TAG, "onStateTransitionEnd - state: " + state);
         }
         for (int i = mListeners.size() - 1; i >= 0; i--) {
@@ -468,7 +480,11 @@
      * Cancels the current animation.
      */
     public void cancelAnimation() {
-        if (DEBUG && mConfig.currentAnimation != null) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logCancelAnimation(
+                    mConfig.currentAnimation != null,
+                    getTrimmedStackTrace("StateManager.cancelAnimation"));
+        } else if (DEBUG && mConfig.currentAnimation != null) {
             Log.d(TAG, "cancelAnimation - with ongoing animation"
                     + ", partial trace:\n" + getTrimmedStackTrace("StateManager.cancelAnimation"));
         }
diff --git a/src/com/android/launcher3/states/RotationHelper.java b/src/com/android/launcher3/states/RotationHelper.java
index fdb37f0..b3bcada 100644
--- a/src/com/android/launcher3/states/RotationHelper.java
+++ b/src/com/android/launcher3/states/RotationHelper.java
@@ -26,8 +26,6 @@
 import static com.android.launcher3.util.window.WindowManagerProxy.MIN_TABLET_WIDTH;
 
 import android.content.Context;
-import android.content.SharedPreferences;
-import android.content.SharedPreferences.OnSharedPreferenceChangeListener;
 import android.os.Handler;
 import android.os.Message;
 
@@ -36,13 +34,14 @@
 
 import com.android.launcher3.BaseActivity;
 import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.LauncherPrefChangeListener;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.util.DisplayController;
 
 /**
  * Utility class to manage launcher rotation
  */
-public class RotationHelper implements OnSharedPreferenceChangeListener,
+public class RotationHelper implements LauncherPrefChangeListener,
         DeviceProfile.OnDeviceProfileChangeListener,
         DisplayController.DisplayInfoChangeListener {
 
@@ -112,7 +111,7 @@
     }
 
     @Override
-    public void onSharedPreferenceChanged(SharedPreferences sharedPreferences, String s) {
+    public void onPrefChanged(String s) {
         if (mDestroyed || mIgnoreAutoRotateSettings) return;
         boolean wasRotationEnabled = mHomeRotationEnabled;
         mHomeRotationEnabled = LauncherPrefs.get(mActivity).get(ALLOW_ROTATION);
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index c59cc81..0b45118 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -35,7 +35,6 @@
 import android.content.ComponentCallbacks;
 import android.content.Context;
 import android.content.Intent;
-import android.content.SharedPreferences;
 import android.content.res.Configuration;
 import android.graphics.Point;
 import android.graphics.Rect;
@@ -51,6 +50,7 @@
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.InvariantDeviceProfile.DeviceType;
+import com.android.launcher3.LauncherPrefChangeListener;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
@@ -116,8 +116,7 @@
     private Info mInfo;
     private boolean mDestroyed = false;
 
-    private SharedPreferences.OnSharedPreferenceChangeListener
-            mTaskbarPinningPreferenceChangeListener;
+    private LauncherPrefChangeListener mTaskbarPinningPreferenceChangeListener;
 
     @VisibleForTesting
     protected DisplayController(Context context) {
@@ -142,19 +141,18 @@
     }
 
     private void attachTaskbarPinningSharedPreferenceChangeListener(Context context) {
-        mTaskbarPinningPreferenceChangeListener =
-                (sharedPreferences, key) -> {
-                    LauncherPrefs prefs = LauncherPrefs.get(mContext);
-                    boolean isTaskbarPinningChanged = TASKBAR_PINNING_KEY.equals(key)
-                            && mInfo.mIsTaskbarPinned != prefs.get(TASKBAR_PINNING);
-                    boolean isTaskbarPinningDesktopModeChanged =
-                            TASKBAR_PINNING_DESKTOP_MODE_KEY.equals(key)
-                                    && mInfo.mIsTaskbarPinnedInDesktopMode != prefs.get(
-                                    TASKBAR_PINNING_IN_DESKTOP_MODE);
-                    if (isTaskbarPinningChanged || isTaskbarPinningDesktopModeChanged) {
-                        notifyConfigChange();
-                    }
-                };
+        mTaskbarPinningPreferenceChangeListener = key -> {
+            LauncherPrefs prefs = LauncherPrefs.get(mContext);
+            boolean isTaskbarPinningChanged = TASKBAR_PINNING_KEY.equals(key)
+                    && mInfo.mIsTaskbarPinned != prefs.get(TASKBAR_PINNING);
+            boolean isTaskbarPinningDesktopModeChanged =
+                    TASKBAR_PINNING_DESKTOP_MODE_KEY.equals(key)
+                            && mInfo.mIsTaskbarPinnedInDesktopMode != prefs.get(
+                            TASKBAR_PINNING_IN_DESKTOP_MODE);
+            if (isTaskbarPinningChanged || isTaskbarPinningDesktopModeChanged) {
+                notifyConfigChange();
+            }
+        };
 
         LauncherPrefs.get(context).addListener(
                 mTaskbarPinningPreferenceChangeListener, TASKBAR_PINNING);
diff --git a/src/com/android/launcher3/util/MainThreadInitializedObject.java b/src/com/android/launcher3/util/MainThreadInitializedObject.java
index a7d5c13..9a70298 100644
--- a/src/com/android/launcher3/util/MainThreadInitializedObject.java
+++ b/src/com/android/launcher3/util/MainThreadInitializedObject.java
@@ -50,7 +50,7 @@
 
     public T get(Context context) {
         Context app = context.getApplicationContext();
-        if (app instanceof SandboxApplication sc) {
+        if (app instanceof ObjectSandbox sc) {
             return sc.getObject(this);
         }
 
@@ -100,7 +100,8 @@
         T get(Context context);
     }
 
-    public interface SandboxApplication {
+    /** Sandbox for isolating {@link MainThreadInitializedObject} instances from Launcher. */
+    public interface ObjectSandbox {
 
         /**
          * Find a cached object from mObjectMap if we have already created one. If not, generate
@@ -116,7 +117,7 @@
         <T extends SafeCloseable> void putObject(MainThreadInitializedObject<T> object, T value);
 
         /**
-         * Returns whether this context should cleanup all objects when its destroyed or leave it
+         * Returns whether this sandbox should cleanup all objects when its destroyed or leave it
          * to the GC.
          * These objects can have listeners attached to the system server and mey not be able to get
          * GCed themselves when running on a device.
@@ -137,7 +138,7 @@
      * Abstract Context which allows custom implementations for
      * {@link MainThreadInitializedObject} providers
      */
-    public static class SandboxContext extends LauncherApplication implements SandboxApplication {
+    public static class SandboxContext extends LauncherApplication implements ObjectSandbox {
 
         private static final String TAG = "SandboxContext";
 
@@ -159,8 +160,8 @@
 
         @Override
         public boolean shouldCleanUpOnDestroy() {
-            return (getBaseContext().getApplicationContext() instanceof SandboxApplication sa)
-                    ? sa.shouldCleanUpOnDestroy() : true;
+            return (getBaseContext().getApplicationContext() instanceof ObjectSandbox os)
+                    ? os.shouldCleanUpOnDestroy() : true;
         }
 
         public void onDestroy() {
diff --git a/src/com/android/launcher3/util/SettingsCache.java b/src/com/android/launcher3/util/SettingsCache.java
index cd6701d..a1ed499 100644
--- a/src/com/android/launcher3/util/SettingsCache.java
+++ b/src/com/android/launcher3/util/SettingsCache.java
@@ -25,14 +25,21 @@
 import android.database.ContentObserver;
 import android.net.Uri;
 import android.os.Handler;
+import android.os.Looper;
 import android.provider.Settings;
 
+import com.android.launcher3.dagger.ApplicationContext;
+import com.android.launcher3.dagger.LauncherAppSingleton;
+import com.android.launcher3.dagger.LauncherBaseAppComponent;
+
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 
+import javax.inject.Inject;
+
 /**
  * ContentObserver over Settings keys that also has a caching layer.
  * Consumers can register for callbacks via {@link #register(Uri, OnChangeListener)} and
@@ -47,6 +54,7 @@
  *
  * Cache will also be updated if a key queried is missing (even if it has no listeners registered).
  */
+@LauncherAppSingleton
 public class SettingsCache extends ContentObserver implements SafeCloseable {
 
     /** Hidden field Settings.Secure.NOTIFICATION_BADGING */
@@ -79,12 +87,14 @@
     /**
      * Singleton instance
      */
-    public static MainThreadInitializedObject<SettingsCache> INSTANCE =
-            new MainThreadInitializedObject<>(SettingsCache::new);
+    public static final DaggerSingletonObject<SettingsCache> INSTANCE =
+            new DaggerSingletonObject<>(LauncherBaseAppComponent::getSettingsCache);
 
-    private SettingsCache(Context context) {
-        super(new Handler());
+    @Inject
+    SettingsCache(@ApplicationContext Context context, DaggerSingletonTracker tracker) {
+        super(new Handler(Looper.getMainLooper()));
         mResolver = context.getContentResolver();
+        ExecutorUtil.executeSyncOnMainOrFail(() -> tracker.addCloseable(this));
     }
 
     @Override
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java b/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java
index 1db3b5a..fba3936 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java
@@ -16,8 +16,8 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.icons.ComponentWithLabelAndIcon;
-import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.BaseIconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 
 /**
@@ -26,8 +26,7 @@
  * (who's implementation is owned by the launcher). This object represents a widget type / class,
  * as opposed to a widget instance, and so should not be confused with {@link LauncherAppWidgetInfo}
  */
-public class LauncherAppWidgetProviderInfo extends AppWidgetProviderInfo
-        implements ComponentWithLabelAndIcon {
+public class LauncherAppWidgetProviderInfo extends AppWidgetProviderInfo implements CachedObject {
 
     public static final String CLS_CUSTOM_WIDGET_PREFIX = "#custom-widget-";
 
@@ -69,6 +68,8 @@
 
     protected boolean mIsMinSizeFulfilled;
 
+    private PackageManager mPM;
+
     public static LauncherAppWidgetProviderInfo fromProviderInfo(Context context,
             AppWidgetProviderInfo info) {
         final LauncherAppWidgetProviderInfo launcherInfo;
@@ -97,6 +98,7 @@
     }
 
     public void initSpans(Context context, InvariantDeviceProfile idp) {
+        mPM = context.getPackageManager();
         int minSpanX = 0;
         int minSpanY = 0;
         int maxSpanX = idp.numColumns;
@@ -104,7 +106,6 @@
         int spanX = 0;
         int spanY = 0;
 
-
         Point cellSize = new Point();
         for (DeviceProfile dp : idp.supportedProfiles) {
             dp.getCellSize(cellSize);
@@ -188,8 +189,9 @@
                 (widgetSize + widgetPadding + cellSpacing) / (cellSize + cellSpacing)));
     }
 
-    public String getLabel(PackageManager packageManager) {
-        return super.loadLabel(packageManager);
+    @Override
+    public CharSequence getLabel() {
+        return super.loadLabel(mPM);
     }
 
     public Point getMinSpans() {
@@ -225,7 +227,7 @@
     }
 
     @Override
-    public Drawable getFullResIcon(IconCache cache) {
+    public Drawable getFullResIcon(BaseIconCache cache) {
         return cache.getFullResIcon(getActivityInfo());
     }
 
diff --git a/src/com/android/launcher3/widget/PendingAddWidgetInfo.java b/src/com/android/launcher3/widget/PendingAddWidgetInfo.java
index a916252..23ab0fb 100644
--- a/src/com/android/launcher3/widget/PendingAddWidgetInfo.java
+++ b/src/com/android/launcher3/widget/PendingAddWidgetInfo.java
@@ -82,8 +82,9 @@
 
     @NonNull
     @Override
-    public LauncherAtom.ItemInfo buildProto(@Nullable CollectionInfo collectionInfo) {
-        LauncherAtom.ItemInfo info = super.buildProto(collectionInfo);
+    public LauncherAtom.ItemInfo buildProto(
+            @Nullable CollectionInfo collectionInfo, Context context) {
+        LauncherAtom.ItemInfo info = super.buildProto(collectionInfo, context);
         return info.toBuilder()
                 .addItemAttributes(LauncherAppWidgetInfo.getAttribute(sourceContainer))
                 .build();
diff --git a/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java b/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
index 5ad9222..82a6883 100644
--- a/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
+++ b/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
@@ -18,7 +18,6 @@
 
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.os.Parcel;
 import android.os.Parcelable;
 
@@ -64,7 +63,7 @@
     }
 
     @Override
-    public String getLabel(PackageManager packageManager) {
+    public CharSequence getLabel() {
         return Utilities.trim(label);
     }
 
diff --git a/src_no_quickstep/com/android/launcher3/util/StateManagerProtoLogProxy.java b/src_no_quickstep/com/android/launcher3/util/StateManagerProtoLogProxy.java
new file mode 100644
index 0000000..34e15f7
--- /dev/null
+++ b/src_no_quickstep/com/android/launcher3/util/StateManagerProtoLogProxy.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util;
+
+/**
+ * Proxy class used for StateManager ProtoLog support.
+ */
+public class StateManagerProtoLogProxy {
+
+    public static void logGoToState(Object fromState, Object toState, String trace) { }
+
+    public static void logCreateAtomicAnimation(Object fromState, Object toState, String trace) { }
+
+    public static void logOnStateTransitionStart(Object state) { }
+
+    public static void logOnStateTransitionEnd(Object state) { }
+
+    public static void logCancelAnimation(boolean animationOngoing, String trace) { }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/AppWidgetsRestoredReceiverTest.kt b/tests/multivalentTests/src/com/android/launcher3/AppWidgetsRestoredReceiverTest.kt
index 21abab4..0e06051 100644
--- a/tests/multivalentTests/src/com/android/launcher3/AppWidgetsRestoredReceiverTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/AppWidgetsRestoredReceiverTest.kt
@@ -29,7 +29,7 @@
 
     @Before
     fun setup() {
-        launcherPrefs = LauncherPrefs(DeviceHelpers.context)
+        launcherPrefs = LauncherPrefs.get(DeviceHelpers.context)
         receiverUnderTest = AppWidgetsRestoredReceiver()
     }
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
new file mode 100644
index 0000000..058ac05
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
@@ -0,0 +1,73 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3
+
+import android.content.Context
+
+/** Emulates Launcher preferences for a test environment. */
+class FakeLauncherPrefs(private val context: Context) : LauncherPrefs() {
+    private val prefsMap = mutableMapOf<String, Any>()
+    private val listeners = mutableSetOf<LauncherPrefChangeListener>()
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <T> get(item: ContextualItem<T>): T {
+        return prefsMap.getOrDefault(item.sharedPrefKey, item.defaultValueFromContext(context)) as T
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <T> get(item: ConstantItem<T>): T {
+        return prefsMap.getOrDefault(item.sharedPrefKey, item.defaultValue) as T
+    }
+
+    override fun put(vararg itemsToValues: Pair<Item, Any>) = putSync(*itemsToValues)
+
+    override fun <T : Any> put(item: Item, value: T) = putSync(item to value)
+
+    override fun putSync(vararg itemsToValues: Pair<Item, Any>) {
+        itemsToValues
+            .map { (i, v) -> i.sharedPrefKey to v }
+            .forEach { (k, v) ->
+                prefsMap[k] = v
+                notifyChange(k)
+            }
+    }
+
+    override fun addListener(listener: LauncherPrefChangeListener, vararg items: Item) {
+        listeners.add(listener)
+    }
+
+    override fun removeListener(listener: LauncherPrefChangeListener, vararg items: Item) {
+        listeners.remove(listener)
+    }
+
+    override fun has(vararg items: Item) = items.all { it.sharedPrefKey in prefsMap }
+
+    override fun remove(vararg items: Item) = removeSync(*items)
+
+    override fun removeSync(vararg items: Item) {
+        items
+            .filter { it.sharedPrefKey in prefsMap }
+            .forEach {
+                prefsMap.remove(it.sharedPrefKey)
+                notifyChange(it.sharedPrefKey)
+            }
+    }
+
+    override fun close() = Unit
+
+    private fun notifyChange(key: String) = listeners.forEach { it.onPrefChanged(key) }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt
new file mode 100644
index 0000000..b8e347c
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt
@@ -0,0 +1,134 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3
+
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+private val TEST_CONSTANT_ITEM = LauncherPrefs.nonRestorableItem("TEST_BOOLEAN_ITEM", false)
+
+private val TEST_CONTEXTUAL_ITEM =
+    ContextualItem(
+        "TEST_CONTEXTUAL_ITEM",
+        true,
+        { false },
+        EncryptionType.ENCRYPTED,
+        Boolean::class.java,
+    )
+
+@RunWith(LauncherMultivalentJUnit::class)
+class FakeLauncherPrefsTest {
+    private val launcherPrefs = FakeLauncherPrefs(getApplicationContext())
+
+    @Test
+    fun testGet_constantItemNotInPrefs_returnsDefaultValue() {
+        assertThat(launcherPrefs.get(TEST_CONSTANT_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testGet_constantItemInPrefs_returnsStoredValue() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        assertThat(launcherPrefs.get(TEST_CONSTANT_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testGet_contextualItemNotInPrefs_returnsDefaultValue() {
+        assertThat(launcherPrefs.get(TEST_CONTEXTUAL_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testGet_contextualItemInPrefs_returnsStoredValue() {
+        launcherPrefs.put(TEST_CONTEXTUAL_ITEM, true)
+        assertThat(launcherPrefs.get(TEST_CONTEXTUAL_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testPut_multipleItems_storesAll() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM to true, TEST_CONTEXTUAL_ITEM to true)
+        assertThat(launcherPrefs.get(TEST_CONSTANT_ITEM)).isTrue()
+        assertThat(launcherPrefs.get(TEST_CONTEXTUAL_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testHas_itemNotInPrefs_returnsFalse() {
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testHas_itemInPrefs_returnsTrue() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testHas_twoItemsWithOneInPrefs_returnsFalse() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM, TEST_CONTEXTUAL_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testHas_twoItemsInPrefs_returnsTrue() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM to true, TEST_CONTEXTUAL_ITEM to true)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM, TEST_CONTEXTUAL_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testRemove_itemInPrefs_removesItem() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        launcherPrefs.remove(TEST_CONSTANT_ITEM)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testRemove_itemsInPrefs_removesItems() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM to true, TEST_CONTEXTUAL_ITEM to true)
+        launcherPrefs.remove(TEST_CONSTANT_ITEM, TEST_CONTEXTUAL_ITEM)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM, TEST_CONTEXTUAL_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testAddListener_changeItemInPrefs_callsListener() {
+        var changedKey: String? = null
+        launcherPrefs.addListener({ changedKey = it }, TEST_CONSTANT_ITEM)
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        assertThat(changedKey).isEqualTo(TEST_CONSTANT_ITEM.sharedPrefKey)
+    }
+
+    @Test
+    fun testAddListener_removeItemFromPrefs_callsListener() {
+        var changedKey: String? = null
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        launcherPrefs.addListener({ changedKey = it }, TEST_CONSTANT_ITEM)
+
+        launcherPrefs.remove(TEST_CONSTANT_ITEM)
+        assertThat(changedKey).isEqualTo(TEST_CONSTANT_ITEM.sharedPrefKey)
+    }
+
+    @Test
+    fun testRemoveListener_changeItemInPrefs_doesNotCallListener() {
+        var changedKey: String? = null
+        val listener = LauncherPrefChangeListener { changedKey = it }
+        launcherPrefs.addListener(listener, TEST_CONSTANT_ITEM)
+
+        launcherPrefs.removeListener(listener)
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        assertThat(changedKey).isNull()
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/LauncherPrefsTest.kt b/tests/multivalentTests/src/com/android/launcher3/LauncherPrefsTest.kt
index b813095..4aeef2e 100644
--- a/tests/multivalentTests/src/com/android/launcher3/LauncherPrefsTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/LauncherPrefsTest.kt
@@ -17,7 +17,6 @@
 
 import android.content.Context
 import android.content.SharedPreferences
-import android.content.SharedPreferences.OnSharedPreferenceChangeListener
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import androidx.test.platform.app.InstrumentationRegistry
@@ -63,7 +62,7 @@
     @Test
     fun addListener_listeningForStringItemUpdates_isCorrectlyNotifiedOfUpdates() {
         val latch = CountDownLatch(1)
-        val listener = OnSharedPreferenceChangeListener { _, _ -> latch.countDown() }
+        val listener = LauncherPrefChangeListener { latch.countDown() }
 
         with(launcherPrefs) {
             putSync(TEST_STRING_ITEM.to(TEST_STRING_ITEM.defaultValue))
@@ -78,7 +77,7 @@
     @Test
     fun removeListener_previouslyListeningForStringItemUpdates_isNoLongerNotifiedOfUpdates() {
         val latch = CountDownLatch(1)
-        val listener = OnSharedPreferenceChangeListener { _, _ -> latch.countDown() }
+        val listener = LauncherPrefChangeListener { latch.countDown() }
 
         with(launcherPrefs) {
             addListener(listener, TEST_STRING_ITEM)
@@ -94,14 +93,14 @@
     @Test
     fun addListenerAndRemoveListener_forMultipleItems_bothWorkProperly() {
         var latch = CountDownLatch(3)
-        val listener = OnSharedPreferenceChangeListener { _, _ -> latch.countDown() }
+        val listener = LauncherPrefChangeListener { latch.countDown() }
 
         with(launcherPrefs) {
             addListener(listener, TEST_INT_ITEM, TEST_STRING_ITEM, TEST_BOOLEAN_ITEM)
             putSync(
                 TEST_INT_ITEM.to(TEST_INT_ITEM.defaultValue + 123),
                 TEST_STRING_ITEM.to(TEST_STRING_ITEM.defaultValue + "abc"),
-                TEST_BOOLEAN_ITEM.to(!TEST_BOOLEAN_ITEM.defaultValue)
+                TEST_BOOLEAN_ITEM.to(!TEST_BOOLEAN_ITEM.defaultValue),
             )
             assertThat(latch.await(WAIT_TIME_IN_SECONDS, TimeUnit.SECONDS)).isTrue()
 
@@ -110,7 +109,7 @@
             putSync(
                 TEST_INT_ITEM.to(TEST_INT_ITEM.defaultValue),
                 TEST_STRING_ITEM.to(TEST_STRING_ITEM.defaultValue),
-                TEST_BOOLEAN_ITEM.to(TEST_BOOLEAN_ITEM.defaultValue)
+                TEST_BOOLEAN_ITEM.to(TEST_BOOLEAN_ITEM.defaultValue),
             )
             remove(TEST_INT_ITEM, TEST_STRING_ITEM, TEST_BOOLEAN_ITEM)
 
@@ -150,7 +149,7 @@
             putSync(
                 TEST_STRING_ITEM.to(TEST_STRING_ITEM.defaultValue),
                 TEST_INT_ITEM.to(TEST_INT_ITEM.defaultValue),
-                TEST_BOOLEAN_ITEM.to(TEST_BOOLEAN_ITEM.defaultValue)
+                TEST_BOOLEAN_ITEM.to(TEST_BOOLEAN_ITEM.defaultValue),
             )
             assertThat(has(TEST_BOOLEAN_ITEM, TEST_INT_ITEM, TEST_STRING_ITEM)).isTrue()
             remove(TEST_STRING_ITEM, TEST_INT_ITEM, TEST_BOOLEAN_ITEM)
@@ -191,7 +190,7 @@
             LauncherPrefs.backedUpItem(
                 TEST_PREF_KEY,
                 TEST_DEFAULT_VALUE,
-                EncryptionType.DEVICE_PROTECTED
+                EncryptionType.DEVICE_PROTECTED,
             )
 
         val bootAwarePrefs: SharedPreferences =
@@ -212,7 +211,7 @@
             LauncherPrefs.backedUpItem(
                 TEST_PREF_KEY,
                 TEST_DEFAULT_VALUE,
-                EncryptionType.DEVICE_PROTECTED
+                EncryptionType.DEVICE_PROTECTED,
             )
 
         val bootAwarePrefs: SharedPreferences =
diff --git a/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt b/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
index d236551..111ffaa 100644
--- a/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
@@ -29,7 +29,6 @@
 import com.android.launcher3.icons.BaseIconFactory
 import com.android.launcher3.icons.FastBitmapDrawable
 import com.android.launcher3.icons.UserBadgeDrawable
-import com.android.launcher3.model.ModelTestRule
 import com.android.launcher3.model.data.FolderInfo
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_ARCHIVED
@@ -45,7 +44,6 @@
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -54,8 +52,6 @@
 @RunWith(AndroidJUnit4::class)
 class PreviewItemManagerTest {
 
-    @get:Rule val modelTestRule = ModelTestRule()
-
     private lateinit var previewItemManager: PreviewItemManager
     private lateinit var context: Context
     private lateinit var folderItems: ArrayList<ItemInfo>
@@ -99,8 +95,8 @@
                 BaseIconFactory(
                     context,
                     context.resources.configuration.densityDpi,
-                    previewItemManager.mIconSize
-                )
+                    previewItemManager.mIconSize,
+                ),
             )
 
         // Set second icon to be non-themed.
@@ -111,8 +107,8 @@
                 BaseIconFactory(
                     context,
                     context.resources.configuration.densityDpi,
-                    previewItemManager.mIconSize
-                )
+                    previewItemManager.mIconSize,
+                ),
             )
 
         // Set third icon to be themed with badge.
@@ -123,8 +119,8 @@
                 BaseIconFactory(
                     context,
                     context.resources.configuration.densityDpi,
-                    previewItemManager.mIconSize
-                )
+                    previewItemManager.mIconSize,
+                ),
             )
         folderApps[2].bitmap = folderApps[2].bitmap.withFlags(profileFlagOp(UserIconInfo.TYPE_WORK))
 
@@ -137,8 +133,8 @@
                 BaseIconFactory(
                     context,
                     context.resources.configuration.densityDpi,
-                    previewItemManager.mIconSize
-                )
+                    previewItemManager.mIconSize,
+                ),
             )
 
         defaultThemedIcons = get(context).get(THEMED_ICONS)
diff --git a/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheTest.java b/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheTest.java
index 519108d..ce00b28 100644
--- a/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheTest.java
@@ -68,7 +68,6 @@
 import com.android.launcher3.util.ApplicationInfoWrapper;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.PackageUserKey;
-import com.android.launcher3.util.RoboApiWrapper;
 
 import com.google.common.truth.Truth;
 
@@ -148,7 +147,6 @@
 
     @Test
     public void launcherActivityInfo_cached_in_memory() {
-        RoboApiWrapper.INSTANCE.initialize();
         ComponentName cn = new ComponentName(TEST_PACKAGE, TEST_ACTIVITY);
         UserHandle user = myUserHandle();
         ComponentKey cacheKey = new ComponentKey(cn, user);
@@ -213,7 +211,6 @@
 
     @Test
     public void item_kept_in_db_if_nothing_changes() {
-        RoboApiWrapper.INSTANCE.initialize();
         ComponentName cn = new ComponentName(TEST_PACKAGE, TEST_ACTIVITY);
         UserHandle user = myUserHandle();
 
@@ -232,7 +229,6 @@
 
     @Test
     public void item_updated_in_db_if_appInfo_changes() {
-        RoboApiWrapper.INSTANCE.initialize();
         ComponentName cn = new ComponentName(TEST_PACKAGE, TEST_ACTIVITY);
         UserHandle user = myUserHandle();
 
@@ -253,7 +249,6 @@
 
     @Test
     public void item_removed_in_db_if_item_removed() {
-        RoboApiWrapper.INSTANCE.initialize();
         ComponentName cn = new ComponentName(TEST_PACKAGE, TEST_ACTIVITY);
         UserHandle user = myUserHandle();
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt b/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt
index 8e54c94..ed9a080 100644
--- a/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt
@@ -18,10 +18,8 @@
 
 import android.content.ComponentName
 import android.content.pm.ApplicationInfo
-import android.content.pm.PackageManager
 import android.database.MatrixCursor
 import android.os.Process.myUserHandle
-import androidx.test.core.app.ApplicationProvider.getApplicationContext
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.icons.cache.BaseIconCache
@@ -49,7 +47,7 @@
     @Mock private lateinit var baseIconCache: BaseIconCache
 
     private var cursor: MatrixCursor? = null
-    private var cachingLogic = CachedObjectCachingLogic<BaseIconCache>(getApplicationContext())
+    private var cachingLogic = CachedObjectCachingLogic
 
     @Before
     fun setup() {
@@ -137,14 +135,13 @@
     }
 }
 
-class TestCachedObject(val cn: ComponentName, val freshnessId: String) :
-    CachedObject<BaseIconCache> {
+class TestCachedObject(val cn: ComponentName, val freshnessId: String) : CachedObject {
 
     override fun getComponent() = cn
 
     override fun getUser() = myUserHandle()
 
-    override fun getLabel(pm: PackageManager?): CharSequence? = null
+    override fun getLabel(): CharSequence? = null
 
     override fun getApplicationInfo(): ApplicationInfo? = null
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt
index 43dc36b..ce04682 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt
@@ -27,7 +27,6 @@
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mockito.times
@@ -44,8 +43,6 @@
 @RunWith(AndroidJUnit4::class)
 class AddWorkspaceItemsTaskTest : AbstractWorkspaceModelTest() {
 
-    @get:Rule val modelTestRule = ModelTestRule()
-
     private lateinit var mDataModelCallbacks: MyCallbacks
 
     private val mWorkspaceItemSpaceFinder: WorkspaceItemSpaceFinder = mock()
@@ -121,18 +118,8 @@
     @Test
     fun givenMultipleItems_whenExecuteTask_thenAddThem() {
         val itemsToAdd =
-            arrayOf(
-                getNewItem(),
-                getExistingItem(),
-                getNewItem(),
-                getNewItem(),
-                getExistingItem(),
-            )
-        givenNewItemSpaces(
-            NewItemSpace(1, 3, 3),
-            NewItemSpace(2, 0, 0),
-            NewItemSpace(2, 0, 1),
-        )
+            arrayOf(getNewItem(), getExistingItem(), getNewItem(), getNewItem(), getExistingItem())
+        givenNewItemSpaces(NewItemSpace(1, 3, 3), NewItemSpace(2, 0, 0), NewItemSpace(2, 0, 1))
         val nonEmptyScreenIds = listOf(0, 1)
 
         val addedItems = testAddItems(nonEmptyScreenIds, *itemsToAdd)
@@ -173,7 +160,7 @@
                 eq(IntArray.wrap(*nonEmptyScreenIds.toIntArray())),
                 eq(IntArray()),
                 eq(1),
-                eq(1)
+                eq(1),
             )
     }
 
@@ -183,7 +170,7 @@
      */
     private fun testAddItems(
         nonEmptyScreenIds: List<Int>,
-        vararg itemsToAdd: WorkspaceItemInfo
+        vararg itemsToAdd: WorkspaceItemInfo,
     ): List<AddedItem> {
         setupWorkspaces(nonEmptyScreenIds)
         val task = newTask(*itemsToAdd)
@@ -220,7 +207,7 @@
     override fun bindAppsAdded(
         newScreens: IntArray?,
         addNotAnimated: ArrayList<ItemInfo>,
-        addAnimated: ArrayList<ItemInfo>
+        addAnimated: ArrayList<ItemInfo>,
     ) {
         addedItems.addAll(addAnimated.map { AddedItem(it, true) })
         addedItems.addAll(addNotAnimated.map { AddedItem(it, false) })
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/AsyncBindingTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/AsyncBindingTest.kt
index dce75b9..ba59253 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/AsyncBindingTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/AsyncBindingTest.kt
@@ -64,8 +64,6 @@
 
     @get:Rule val setFlagsRule = SetFlagsRule()
 
-    @get:Rule val modelTestRule = ModelTestRule()
-
     @Spy private var callbacks = MyCallbacks()
     @Mock private lateinit var itemInflater: ItemInflater<*>
 
@@ -138,7 +136,7 @@
     @Test
     fun test_bind_sync_partially_inflates_on_background() {
         modelHelper.loadModelSync()
-        assertTrue(modelHelper.model.isModelLoaded)
+        assertTrue(modelHelper.model.isModelLoaded())
         callbacks.inflater = itemInflater
 
         val firstPageBindIds = IntSet()
@@ -203,7 +201,7 @@
             pendingTasks: RunnableList,
             onCompleteSignal: RunnableList,
             workspaceItemCount: Int,
-            isBindSync: Boolean
+            isBindSync: Boolean,
         ) {
             this.pendingTasks = pendingTasks
             this.onCompleteSignal = onCompleteSignal
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java b/tests/multivalentTests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
index 535080a..600af42 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/model/CacheDataUpdatedTaskTest.java
@@ -64,11 +64,7 @@
 @RunWith(AndroidJUnit4.class)
 public class CacheDataUpdatedTaskTest {
 
-    @Rule(order = 0)
-    public TestRule testStabilityRule = new TestStabilityRule();
-
-    @Rule(order = 1)
-    public ModelTestRule mModelTestRule = new ModelTestRule();
+    @Rule public TestRule testStabilityRule = new TestStabilityRule();
 
     private static final String PENDING_APP_1 = TEST_PACKAGE + ".pending1";
     private static final String PENDING_APP_2 = TEST_PACKAGE + ".pending2";
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java b/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
index e14e145..1e2431f 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
@@ -39,7 +39,6 @@
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -50,8 +49,6 @@
 @RunWith(AndroidJUnit4.class)
 public class DefaultLayoutProviderTest {
 
-    @Rule public ModelTestRule rule = new ModelTestRule();
-
     private LauncherModelHelper mModelHelper;
     private LauncherModelHelper.SandboxModelContext mTargetContext;
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/FirstScreenBroadcastHelperTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/FirstScreenBroadcastHelperTest.kt
index d2d9512..9cc380e 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/FirstScreenBroadcastHelperTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/FirstScreenBroadcastHelperTest.kt
@@ -34,7 +34,6 @@
 import com.android.launcher3.util.PackageManagerHelper
 import com.android.launcher3.util.PackageUserKey
 import junit.framework.Assert.assertEquals
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
@@ -46,8 +45,6 @@
 @RunWith(AndroidJUnit4::class)
 class FirstScreenBroadcastHelperTest {
 
-    @get:Rule val modelTestRule = ModelTestRule()
-
     private val context = spy(InstrumentationRegistry.getInstrumentation().targetContext)
     private val mockPmHelper = mock<PackageManagerHelper>()
     private val expectedAppPackage = "appPackageExpected"
@@ -70,7 +67,7 @@
                 container = CONTAINER_HOTSEAT
                 intent = expectedIntent
             },
-            LauncherAppWidgetInfo().apply { providerName = expectedComponentName }
+            LauncherAppWidgetInfo().apply { providerName = expectedComponentName },
         )
 
     @Test
@@ -89,7 +86,7 @@
         val sessionInfoMap: HashMap<PackageUserKey, SessionInfo> =
             hashMapOf(
                 PackageUserKey(unexpectedAppPackage, UserHandle(0)) to sessionInfoExpected,
-                PackageUserKey(expectedAppPackage, UserHandle(0)) to sessionInfoUnexpected
+                PackageUserKey(expectedAppPackage, UserHandle(0)) to sessionInfoUnexpected,
             )
 
         // When
@@ -98,7 +95,7 @@
                 packageManagerHelper = mockPmHelper,
                 firstScreenItems = firstScreenItems,
                 userKeyToSessionMap = sessionInfoMap,
-                allWidgets = listOf()
+                allWidgets = listOf(),
             )
 
         // Then
@@ -108,7 +105,7 @@
                     installerPackage = expectedInstallerPackage,
                     pendingWorkspaceItems = mutableSetOf(expectedAppPackage),
                     pendingHotseatItems = mutableSetOf(expectedAppPackage),
-                    pendingWidgetItems = mutableSetOf(expectedAppPackage)
+                    pendingWidgetItems = mutableSetOf(expectedAppPackage),
                 )
             )
 
@@ -133,7 +130,7 @@
                             providerName = expectedComponentName
                             screenId = 0
                         }
-                    )
+                    ),
             )
 
         // Then
@@ -143,7 +140,7 @@
                     installerPackage = expectedInstallerPackage,
                     installedHotseatItems = mutableSetOf(expectedAppPackage),
                     installedWorkspaceItems = mutableSetOf(expectedAppPackage),
-                    firstScreenInstalledWidgets = mutableSetOf(expectedAppPackage)
+                    firstScreenInstalledWidgets = mutableSetOf(expectedAppPackage),
                 )
             )
         assertEquals(expectedResult, actualResult)
@@ -178,8 +175,8 @@
                         LauncherAppWidgetInfo().apply {
                             providerName = unexpectedComponentName
                             screenId = 0
-                        }
-                    )
+                        },
+                    ),
             )
 
         // Then
@@ -190,7 +187,7 @@
                     installedHotseatItems = mutableSetOf(),
                     installedWorkspaceItems = mutableSetOf(),
                     firstScreenInstalledWidgets = mutableSetOf(expectedAppPackage),
-                    secondaryScreenInstalledWidgets = mutableSetOf(expectedAppPackage2)
+                    secondaryScreenInstalledWidgets = mutableSetOf(expectedAppPackage2),
                 )
             )
         assertEquals(expectedResult, actualResult)
@@ -224,7 +221,7 @@
                 packageManagerHelper = mockPmHelper,
                 firstScreenItems = firstScreenItems,
                 userKeyToSessionMap = sessionInfoMap,
-                allWidgets = listOf()
+                allWidgets = listOf(),
             )
 
         // Then
@@ -232,7 +229,7 @@
             listOf(
                 FirstScreenBroadcastModel(
                     installerPackage = expectedInstallerPackage,
-                    pendingCollectionItems = mutableSetOf(expectedAppPackage)
+                    pendingCollectionItems = mutableSetOf(expectedAppPackage),
                 )
             )
         assertEquals(expectedResult, actualResult)
@@ -259,7 +256,7 @@
                 firstScreenInstalledWidgets =
                     mutableSetOf<String>().apply { repeat(20) { add(it.toString()) } },
                 secondaryScreenInstalledWidgets =
-                    mutableSetOf<String>().apply { repeat(20) { add(it.toString()) } }
+                    mutableSetOf<String>().apply { repeat(20) { add(it.toString()) } },
             )
 
         // When
@@ -334,7 +331,7 @@
                     installedWorkspaceItems = mutableSetOf("installedWorkspaceItems"),
                     installedHotseatItems = mutableSetOf("installedHotseatItems"),
                     firstScreenInstalledWidgets = mutableSetOf("firstScreenInstalledWidgetItems"),
-                    secondaryScreenInstalledWidgets = mutableSetOf("secondaryInstalledWidgetItems")
+                    secondaryScreenInstalledWidgets = mutableSetOf("secondaryInstalledWidgetItems"),
                 )
             )
         val expectedPendingIntent =
@@ -342,7 +339,7 @@
                 context,
                 0 /* requestCode */,
                 Intent(),
-                PendingIntent.FLAG_ONE_SHOT or PendingIntent.FLAG_IMMUTABLE
+                PendingIntent.FLAG_ONE_SHOT or PendingIntent.FLAG_IMMUTABLE,
             )
 
         // When
@@ -354,40 +351,40 @@
 
         assertEquals(
             "com.android.launcher3.action.FIRST_SCREEN_ACTIVE_INSTALLS",
-            argumentCaptor.value.action
+            argumentCaptor.value.action,
         )
         assertEquals(expectedInstallerPackage, argumentCaptor.value.`package`)
         assertEquals(
             expectedPendingIntent,
-            argumentCaptor.value.getParcelableExtra("verificationToken")
+            argumentCaptor.value.getParcelableExtra("verificationToken"),
         )
         assertEquals(
             arrayListOf("pendingCollectionItem"),
-            argumentCaptor.value.getStringArrayListExtra("folderItem")
+            argumentCaptor.value.getStringArrayListExtra("folderItem"),
         )
         assertEquals(
             arrayListOf("pendingWorkspaceItem"),
-            argumentCaptor.value.getStringArrayListExtra("workspaceItem")
+            argumentCaptor.value.getStringArrayListExtra("workspaceItem"),
         )
         assertEquals(
             arrayListOf("pendingHotseatItems"),
-            argumentCaptor.value.getStringArrayListExtra("hotseatItem")
+            argumentCaptor.value.getStringArrayListExtra("hotseatItem"),
         )
         assertEquals(
             arrayListOf("pendingWidgetItems"),
-            argumentCaptor.value.getStringArrayListExtra("widgetItem")
+            argumentCaptor.value.getStringArrayListExtra("widgetItem"),
         )
         assertEquals(
             arrayListOf("installedWorkspaceItems"),
-            argumentCaptor.value.getStringArrayListExtra("workspaceInstalledItems")
+            argumentCaptor.value.getStringArrayListExtra("workspaceInstalledItems"),
         )
         assertEquals(
             arrayListOf("installedHotseatItems"),
-            argumentCaptor.value.getStringArrayListExtra("hotseatInstalledItems")
+            argumentCaptor.value.getStringArrayListExtra("hotseatInstalledItems"),
         )
         assertEquals(
             arrayListOf("firstScreenInstalledWidgetItems", "secondaryInstalledWidgetItems"),
-            argumentCaptor.value.getStringArrayListExtra("widgetInstalledItems")
+            argumentCaptor.value.getStringArrayListExtra("widgetInstalledItems"),
         )
     }
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
index 4ca47e3..e8f778f 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
@@ -30,7 +30,6 @@
 import com.google.common.truth.Truth.assertWithMessage
 import org.junit.After
 import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -39,8 +38,6 @@
 @RunWith(AndroidJUnit4::class)
 class FolderIconLoadTest {
 
-    @get:Rule(order = 0) val modelTestRule = ModelTestRule()
-
     private lateinit var modelHelper: LauncherModelHelper
 
     private val uniqueActivities =
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/LoaderCursorTest.java b/tests/multivalentTests/src/com/android/launcher3/model/LoaderCursorTest.java
index ac911b3..b4945d7 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/LoaderCursorTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/model/LoaderCursorTest.java
@@ -67,7 +67,6 @@
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -78,8 +77,6 @@
 @RunWith(AndroidJUnit4.class)
 public class LoaderCursorTest {
 
-    @Rule public ModelTestRule rule = new ModelTestRule();
-
     private LauncherModelHelper mModelHelper;
     private LauncherAppState mApp;
     private PackageManagerHelper mPmHelper;
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/ModelTestRule.kt b/tests/multivalentTests/src/com/android/launcher3/model/ModelTestRule.kt
deleted file mode 100644
index ad2c2a4..0000000
--- a/tests/multivalentTests/src/com/android/launcher3/model/ModelTestRule.kt
+++ /dev/null
@@ -1,27 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.model
-
-import com.android.launcher3.util.RoboApiWrapper
-import org.junit.rules.TestWatcher
-import org.junit.runner.Description
-
-class ModelTestRule : TestWatcher() {
-    override fun starting(description: Description?) {
-        RoboApiWrapper.initialize()
-    }
-}
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java b/tests/multivalentTests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
index a0d9da9..0f1fc00 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/model/PackageInstallStateChangedTaskTest.java
@@ -37,7 +37,6 @@
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -48,8 +47,6 @@
 @RunWith(AndroidJUnit4.class)
 public class PackageInstallStateChangedTaskTest {
 
-    @Rule public ModelTestRule mModelTestRule = new ModelTestRule();
-
     private static final String PENDING_APP_1 = TEST_PACKAGE + ".pending1";
     private static final String PENDING_APP_2 = TEST_PACKAGE + ".pending2";
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt
index c7abce6..ed8b397 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt
@@ -58,7 +58,6 @@
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
@@ -77,8 +76,6 @@
 @RunWith(AndroidJUnit4::class)
 class WorkspaceItemProcessorTest {
 
-    @get:Rule val modelTestRule = ModelTestRule()
-
     @Mock private lateinit var mockIconRequestInfo: IconRequestInfo<WorkspaceItemInfo>
     @Mock private lateinit var mockWorkspaceInfo: WorkspaceItemInfo
     @Mock private lateinit var mockBgDataModel: BgDataModel
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemSpaceFinderTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemSpaceFinderTest.kt
index ae8e966..dd03eee 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemSpaceFinderTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemSpaceFinderTest.kt
@@ -21,7 +21,6 @@
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -30,8 +29,6 @@
 @RunWith(AndroidJUnit4::class)
 class WorkspaceItemSpaceFinderTest : AbstractWorkspaceModelTest() {
 
-    @get:Rule val modelTestRule = ModelTestRule()
-
     private val mItemSpaceFinder = WorkspaceItemSpaceFinder()
 
     @Before
@@ -52,7 +49,7 @@
                 mExistingScreens,
                 mNewScreens,
                 spanX,
-                spanY
+                spanY,
             )
             .let { NewItemSpace.fromIntArray(it) }
 
@@ -62,7 +59,7 @@
                     newItemSpace.cellX,
                     newItemSpace.cellY,
                     spanX,
-                    spanY
+                    spanY,
                 )
             )
             .isTrue()
@@ -171,7 +168,7 @@
             screen0 = listOf(Rect(2, 0, 5, 2)),
             screen1 = fullScreenSpaces, // full screens are skipped
             screen2 = fullScreenSpaces, // full screens are skipped
-            screen3 = emptyScreenSpaces
+            screen3 = emptyScreenSpaces,
         )
 
         val spaceFound = findSpace(3, 1)
diff --git a/tests/multivalentTests/src/com/android/launcher3/pm/InstallSessionTrackerTest.kt b/tests/multivalentTests/src/com/android/launcher3/pm/InstallSessionTrackerTest.kt
index d860710..15a9964 100644
--- a/tests/multivalentTests/src/com/android/launcher3/pm/InstallSessionTrackerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/pm/InstallSessionTrackerTest.kt
@@ -26,7 +26,6 @@
 import androidx.test.filters.SdkSuppress
 import androidx.test.filters.SmallTest
 import com.android.launcher3.Flags.FLAG_ENABLE_SUPPORT_FOR_ARCHIVING
-import com.android.launcher3.model.ModelTestRule
 import com.android.launcher3.util.Executors.MODEL_EXECUTOR
 import com.android.launcher3.util.LauncherModelHelper
 import com.android.launcher3.util.PackageUserKey
@@ -45,9 +44,7 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class InstallSessionTrackerTest {
-    @get:Rule(order = 0) val setFlagsRule = SetFlagsRule()
-
-    @get:Rule(order = 1) val modelTestRule = ModelTestRule()
+    @get:Rule val setFlagsRule = SetFlagsRule()
 
     private val mockInstallSessionHelper: InstallSessionHelper = mock()
     private val mockCallback: InstallSessionTracker.Callback = mock()
@@ -67,7 +64,7 @@
                 mockInstallSessionHelper,
                 mockCallback,
                 mockPackageInstaller,
-                launcherApps
+                launcherApps,
             )
     }
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/pm/UserCacheTest.kt b/tests/multivalentTests/src/com/android/launcher3/pm/UserCacheTest.kt
index 482dced..5f08c31 100644
--- a/tests/multivalentTests/src/com/android/launcher3/pm/UserCacheTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/pm/UserCacheTest.kt
@@ -20,7 +20,6 @@
 import android.os.UserHandle
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
-import com.android.launcher3.model.ModelTestRule
 import com.android.launcher3.util.Executors.MODEL_EXECUTOR
 import com.android.launcher3.util.LauncherModelHelper
 import com.android.launcher3.util.TestUtil
@@ -28,15 +27,12 @@
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
 @RunWith(AndroidJUnit4::class)
 class UserCacheTest {
 
-    @get:Rule val modelTestRule = ModelTestRule()
-
     private val launcherModelHelper = LauncherModelHelper()
     private val sandboxContext = launcherModelHelper.sandboxContext
     private lateinit var userCache: UserCache
diff --git a/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index b3675a6..d0c168a 100644
--- a/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -101,7 +101,7 @@
         mMockController = Mockito.mock(ModelDbController.class);
         mMockDb = mock(SQLiteDatabase.class);
         mMockCursor = mock(Cursor.class);
-        mPrefs = new LauncherPrefs(mContext);
+        mPrefs = LauncherPrefs.get(mContext);
         mMockRestoreEventLogger = mock(LauncherRestoreEventLogger.class);
     }
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
index 748d376..09b9a3b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -31,6 +31,7 @@
 
 import android.content.ContentProvider;
 import android.content.ContentResolver;
+import android.content.Context;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageInstaller.SessionParams;
 import android.content.pm.PackageManager;
@@ -250,15 +251,16 @@
         private final File mDbDir;
 
         public SandboxModelContext() {
-            super(ApplicationProvider.getApplicationContext());
+            this(ApplicationProvider.getApplicationContext());
+        }
+
+        public SandboxModelContext(Context context) {
+            super(context);
 
             // System settings cache content provider. Ensure that they are statically initialized
-            Settings.Secure.getString(
-                    ApplicationProvider.getApplicationContext().getContentResolver(), "test");
-            Settings.System.getString(
-                    ApplicationProvider.getApplicationContext().getContentResolver(), "test");
-            Settings.Global.getString(
-                    ApplicationProvider.getApplicationContext().getContentResolver(), "test");
+            Settings.Secure.getString(context.getContentResolver(), "test");
+            Settings.System.getString(context.getContentResolver(), "test");
+            Settings.Global.getString(context.getContentResolver(), "test");
 
             mPm = spy(getBaseContext().getPackageManager());
             mDbDir = new File(getCacheDir(), UUID.randomUUID().toString());
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/SandboxApplication.kt b/tests/multivalentTests/src/com/android/launcher3/util/SandboxApplication.kt
new file mode 100644
index 0000000..4f9b8c7
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/SandboxApplication.kt
@@ -0,0 +1,166 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util
+
+import android.content.Context
+import android.content.ContextParams
+import android.content.ContextWrapper
+import android.content.pm.ApplicationInfo
+import android.content.res.Configuration
+import android.os.Bundle
+import android.os.IBinder
+import android.os.UserHandle
+import android.view.Display
+import androidx.test.core.app.ApplicationProvider
+import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext
+import com.android.launcher3.util.MainThreadInitializedObject.ObjectSandbox
+import org.junit.Rule
+import org.junit.rules.ExternalResource
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * Sandbox application where created [Context] instances are still sandboxed within it.
+ *
+ * Tests can declare this application as a [Rule], so that it is set up and destroyed automatically.
+ * Alternatively, they can call [init] and [onDestroy] directly. Either way, these need to be called
+ * for it to work and avoid leaks from created singletons.
+ *
+ * The create [Context] APIs construct a `ContextImpl`, which resets the application to the true
+ * application, thus leaving the sandbox. This implementation wraps the created contexts to
+ * propagate this application (see [SandboxApplicationWrapper]).
+ */
+class SandboxApplication private constructor(private val base: SandboxApplicationWrapper) :
+    SandboxModelContext(base), TestRule {
+
+    constructor(
+        base: Context = ApplicationProvider.getApplicationContext()
+    ) : this(SandboxApplicationWrapper(base))
+
+    /**
+     * Initializes the sandbox application propagation logic.
+     *
+     * This function either needs to be called manually or automatically through using [Rule].
+     */
+    fun init() {
+        base.app = this@SandboxApplication
+    }
+
+    /** Returns `this` if [init] was called, otherwise crashes the test. */
+    override fun getApplicationContext(): Context = base.applicationContext
+
+    override fun shouldCleanUpOnDestroy(): Boolean {
+        // Defer to the true application to decide whether to clean up. For instance, we do not want
+        // to cleanup under Robolectric.
+        val app = ApplicationProvider.getApplicationContext<Context>()
+        return if (app is ObjectSandbox) app.shouldCleanUpOnDestroy() else true
+    }
+
+    override fun apply(statement: Statement, description: Description): Statement {
+        return object : ExternalResource() {
+                override fun before() {
+                    base.app = this@SandboxApplication
+                }
+
+                override fun after() = onDestroy()
+            }
+            .apply(statement, description)
+    }
+}
+
+private class SandboxApplicationWrapper(base: Context, var app: Context? = null) :
+    ContextWrapper(base) {
+
+    override fun getApplicationContext(): Context {
+        return checkNotNull(app) { "SandboxApplication accessed before #init() was called." }
+    }
+
+    override fun createPackageContext(packageName: String?, flags: Int): Context {
+        return SandboxApplicationWrapper(super.createPackageContext(packageName, flags), app)
+    }
+
+    override fun createPackageContextAsUser(
+        packageName: String,
+        flags: Int,
+        user: UserHandle,
+    ): Context {
+        return SandboxApplicationWrapper(
+            super.createPackageContextAsUser(packageName, flags, user),
+            app,
+        )
+    }
+
+    override fun createContextAsUser(user: UserHandle, flags: Int): Context {
+        return SandboxApplicationWrapper(super.createContextAsUser(user, flags), app)
+    }
+
+    override fun createApplicationContext(application: ApplicationInfo?, flags: Int): Context {
+        return SandboxApplicationWrapper(super.createApplicationContext(application, flags), app)
+    }
+
+    override fun createContextForSdkInSandbox(sdkInfo: ApplicationInfo, flags: Int): Context {
+        return SandboxApplicationWrapper(super.createContextForSdkInSandbox(sdkInfo, flags), app)
+    }
+
+    override fun createContextForSplit(splitName: String?): Context {
+        return SandboxApplicationWrapper(super.createContextForSplit(splitName), app)
+    }
+
+    override fun createConfigurationContext(overrideConfiguration: Configuration): Context {
+        return SandboxApplicationWrapper(
+            super.createConfigurationContext(overrideConfiguration),
+            app,
+        )
+    }
+
+    override fun createDisplayContext(display: Display): Context {
+        return SandboxApplicationWrapper(super.createDisplayContext(display), app)
+    }
+
+    override fun createDeviceContext(deviceId: Int): Context {
+        return SandboxApplicationWrapper(super.createDeviceContext(deviceId), app)
+    }
+
+    override fun createWindowContext(type: Int, options: Bundle?): Context {
+        return SandboxApplicationWrapper(super.createWindowContext(type, options), app)
+    }
+
+    override fun createWindowContext(display: Display, type: Int, options: Bundle?): Context {
+        return SandboxApplicationWrapper(super.createWindowContext(display, type, options), app)
+    }
+
+    override fun createContext(contextParams: ContextParams): Context {
+        return SandboxApplicationWrapper(super.createContext(contextParams), app)
+    }
+
+    override fun createAttributionContext(attributionTag: String?): Context {
+        return SandboxApplicationWrapper(super.createAttributionContext(attributionTag), app)
+    }
+
+    override fun createCredentialProtectedStorageContext(): Context {
+        return SandboxApplicationWrapper(super.createCredentialProtectedStorageContext(), app)
+    }
+
+    override fun createDeviceProtectedStorageContext(): Context {
+        return SandboxApplicationWrapper(super.createDeviceProtectedStorageContext(), app)
+    }
+
+    override fun createTokenContext(token: IBinder, display: Display): Context {
+        return SandboxApplicationWrapper(super.createTokenContext(token, display), app)
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/SandboxApplicationTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/SandboxApplicationTest.kt
new file mode 100644
index 0000000..d87a406
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/SandboxApplicationTest.kt
@@ -0,0 +1,78 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util
+
+import android.content.Context
+import android.hardware.display.DisplayManager
+import android.view.Display
+import android.view.Display.DEFAULT_DISPLAY
+import android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY
+import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(LauncherMultivalentJUnit::class)
+class SandboxApplicationTest {
+    @get:Rule val app = SandboxApplication()
+
+    private val display: Display
+        get() {
+            return checkNotNull(app.getSystemService(DisplayManager::class.java))
+                .getDisplay(DEFAULT_DISPLAY)
+        }
+
+    @Test
+    fun testCreateDisplayContext_isSandboxed() {
+        val displayContext = app.createDisplayContext(display)
+        assertThat(displayContext.applicationContext).isEqualTo(app)
+    }
+
+    @Test
+    fun testCreateWindowContext_fromSandboxedDisplayContext_isSandboxed() {
+        val displayContext = app.createDisplayContext(display)
+        val nestedContext = displayContext.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
+        assertThat(nestedContext.applicationContext).isEqualTo(app)
+    }
+
+    @Test(expected = IllegalStateException::class)
+    fun testGetApplicationContext_beforeManualInit_throwsException() {
+        val manualApp = SandboxApplication()
+        assertThat(manualApp.applicationContext).isEqualTo(manualApp)
+    }
+
+    @Test
+    fun testGetApplicationContext_afterManualInit_isApplication() {
+        SandboxApplication().run {
+            init()
+            assertThat(applicationContext).isEqualTo(this)
+            onDestroy()
+        }
+    }
+
+    @Test
+    fun testGetObject_objectCreatesDisplayContext_isSandboxed() {
+        class TestSingleton(context: Context) : SafeCloseable {
+            override fun close() = Unit
+
+            val displayContext = context.createDisplayContext(display)
+        }
+
+        val displayContext = MainThreadInitializedObject { TestSingleton(it) }[app].displayContext
+        assertThat(displayContext.applicationContext).isEqualTo(app)
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/TestUtil.java b/tests/multivalentTests/src/com/android/launcher3/util/TestUtil.java
index 64035da..ce682f1 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/TestUtil.java
+++ b/tests/multivalentTests/src/com/android/launcher3/util/TestUtil.java
@@ -15,14 +15,12 @@
  */
 package com.android.launcher3.util;
 
-import static android.util.Base64.NO_PADDING;
-import static android.util.Base64.NO_WRAP;
-
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
-import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_KEY;
 import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_LABEL;
 import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_TAG;
+import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_PROVIDER_KEY;
+import static com.android.launcher3.LauncherSettings.Settings.createBlobProviderKey;
 
 import static org.junit.Assert.assertTrue;
 
@@ -42,7 +40,6 @@
 import android.os.UserHandle;
 import android.provider.Settings;
 import android.system.OsConstants;
-import android.util.Base64;
 import android.util.Log;
 
 import androidx.test.uiautomator.UiDevice;
@@ -169,13 +166,12 @@
             session.commit(AsyncTask.THREAD_POOL_EXECUTOR, i -> wait.countDown());
         }
 
-        String key = Base64.encodeToString(digest, NO_WRAP | NO_PADDING);
-
         grantWriteSecurePermission();
-        Settings.Secure.putString(context.getContentResolver(), LAYOUT_DIGEST_KEY, key);
+        Settings.Secure.putString(
+                context.getContentResolver(), LAYOUT_PROVIDER_KEY, createBlobProviderKey(digest));
         wait.await();
         return () ->
-            Settings.Secure.putString(context.getContentResolver(), LAYOUT_DIGEST_KEY, null);
+            Settings.Secure.putString(context.getContentResolver(), LAYOUT_PROVIDER_KEY, null);
     }
 
     /**
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt
index 0a3035a..af2c378 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt
@@ -17,7 +17,6 @@
 package com.android.launcher3.widget.custom
 
 import android.content.ComponentName
-import android.content.pm.PackageManager
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
@@ -25,7 +24,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mockito.mock
 
 @MediumTest
 @RunWith(AndroidJUnit4::class)
@@ -47,7 +45,7 @@
     @Test
     fun get_label() {
         underTest.label = "  TEST_LABEL"
-        assertEquals(LABEL_NAME, underTest.getLabel(mock(PackageManager::class.java)))
+        assertEquals(LABEL_NAME, underTest.getLabel())
     }
 
     companion object {
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/model/WidgetsListBaseEntriesBuilderTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/model/WidgetsListBaseEntriesBuilderTest.kt
index 5df7caa..063ab32 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/model/WidgetsListBaseEntriesBuilderTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/model/WidgetsListBaseEntriesBuilderTest.kt
@@ -26,8 +26,8 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherAppState
-import com.android.launcher3.icons.ComponentWithLabel
 import com.android.launcher3.icons.IconCache
+import com.android.launcher3.icons.cache.CachedObject
 import com.android.launcher3.model.WidgetItem
 import com.android.launcher3.model.data.PackageItemInfo
 import com.android.launcher3.util.ActivityContextWrapper
@@ -66,11 +66,11 @@
         testInvariantProfile = LauncherAppState.getIDP(context)
 
         doAnswer { invocation: InvocationOnMock ->
-                val componentWithLabel = invocation.getArgument<Any>(0) as ComponentWithLabel
+                val componentWithLabel = invocation.getArgument<Any>(0) as CachedObject
                 componentWithLabel.getComponent().shortClassName
             }
             .`when`(iconCache)
-            .getTitleNoCache(any<ComponentWithLabel>())
+            .getTitleNoCache(any<CachedObject>())
         underTest = WidgetsListBaseEntriesBuilder(context)
 
         allWidgets =
@@ -79,14 +79,14 @@
                 packageItemInfoWithTitle(APP_1_PACKAGE_NAME, APP_1_PACKAGE_TITLE) to
                     listOf(
                         createWidgetItem(APP_1_PACKAGE_NAME, APP_1_PROVIDER_1_CLASS_NAME),
-                        createWidgetItem(APP_1_PACKAGE_NAME, APP_1_PROVIDER_2_CLASS_NAME)
+                        createWidgetItem(APP_1_PACKAGE_NAME, APP_1_PROVIDER_2_CLASS_NAME),
                     ),
                 // app 2
                 packageItemInfoWithTitle(APP_2_PACKAGE_NAME, APP_2_PACKAGE_TITLE) to
                     listOf(createWidgetItem(APP_2_PACKAGE_NAME, APP_2_PROVIDER_1_CLASS_NAME)),
                 // app 3
                 packageItemInfoWithTitle(APP_3_PACKAGE_NAME, APP_3_PACKAGE_TITLE) to
-                    listOf(createWidgetItem(APP_3_PACKAGE_NAME, APP_3_PROVIDER_1_CLASS_NAME))
+                    listOf(createWidgetItem(APP_3_PACKAGE_NAME, APP_3_PROVIDER_1_CLASS_NAME)),
             )
     }
 
@@ -96,7 +96,7 @@
             listOf(
                 APP_1_EXPECTED_SECTION_NAME to 2,
                 APP_2_EXPECTED_SECTION_NAME to 1,
-                APP_3_EXPECTED_SECTION_NAME to 1
+                APP_3_EXPECTED_SECTION_NAME to 1,
             )
 
         val entries = underTest.build(allWidgets)
@@ -122,7 +122,7 @@
         val expectedWidgetsCountBySection =
             listOf(
                 APP_1_EXPECTED_SECTION_NAME to 1, // one widget filtered out
-                APP_3_EXPECTED_SECTION_NAME to 1
+                APP_3_EXPECTED_SECTION_NAME to 1,
             )
 
         val entries =
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
index d4e061a..c9b6d4f 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -42,8 +42,8 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
@@ -87,7 +87,7 @@
         mTestProfile.numColumns = 5;
 
         doAnswer(invocation -> {
-            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            CachedObject componentWithLabel = invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
         mViewHolderBinder = new WidgetsListHeaderViewHolderBinder(
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
index e1cc010..0d9464a 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
@@ -45,8 +45,8 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
@@ -92,7 +92,7 @@
         mTestProfile.numColumns = 5;
 
         doAnswer(invocation -> {
-            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            CachedObject componentWithLabel = invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt
index 1822639..1da74cb 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt
@@ -27,8 +27,8 @@
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.LauncherSettings
-import com.android.launcher3.icons.ComponentWithLabel
 import com.android.launcher3.icons.IconCache
+import com.android.launcher3.icons.cache.CachedObject
 import com.android.launcher3.model.WidgetItem
 import com.android.launcher3.model.data.PackageItemInfo
 import com.android.launcher3.util.ActivityContextWrapper
@@ -81,11 +81,11 @@
         testInvariantProfile = LauncherAppState.getIDP(context)
 
         doAnswer { invocation: InvocationOnMock ->
-                val componentWithLabel = invocation.getArgument<Any>(0) as ComponentWithLabel
+                val componentWithLabel = invocation.getArgument<Any>(0) as CachedObject
                 componentWithLabel.getComponent().shortClassName
             }
             .`when`(iconCache)
-            .getTitleNoCache(any<ComponentWithLabel>())
+            .getTitleNoCache(any<CachedObject>())
 
         appWidgetItem = createWidgetItem()
     }
@@ -113,8 +113,8 @@
             listOf(
                 PendingAddWidgetInfo(
                     appWidgetItem.widgetInfo,
-                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
-                ),
+                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION,
+                )
             )
         underTest.setWidgetRecommendations(recommendations)
 
@@ -133,8 +133,8 @@
             listOf(
                 PendingAddWidgetInfo(
                     appWidgetItem.widgetInfo,
-                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
-                ),
+                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION,
+                )
             )
         underTest.setWidgetRecommendations(recommendations)
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
index 7552619..6088c8e 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
@@ -33,8 +33,8 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
@@ -81,7 +81,7 @@
         mTestProfile.numColumns = 5;
 
         doAnswer(invocation -> {
-            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            CachedObject componentWithLabel = invocation.getArgument(0);
             return mWidgetsToLabels.get(componentWithLabel.getComponent());
         }).when(mIconCache).getTitleNoCache(any());
     }
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt
index e59e211..deec67a 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt
@@ -27,8 +27,8 @@
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
-import com.android.launcher3.icons.ComponentWithLabel
 import com.android.launcher3.icons.IconCache
+import com.android.launcher3.icons.cache.CachedObject
 import com.android.launcher3.model.WidgetItem
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.model.data.PackageItemInfo
@@ -86,11 +86,11 @@
         testInvariantProfile = LauncherAppState.getIDP(context)
 
         doAnswer { invocation: InvocationOnMock ->
-                val componentWithLabel = invocation.getArgument<Any>(0) as ComponentWithLabel
+                val componentWithLabel = invocation.getArgument<Any>(0) as CachedObject
                 componentWithLabel.getComponent().shortClassName
             }
             .`when`(iconCache)
-            .getTitleNoCache(any<ComponentWithLabel>())
+            .getTitleNoCache(any<CachedObject>())
 
         app1PackageItemInfo = packageItemInfoWithTitle(APP_1_PACKAGE_NAME, APP_1_PACKAGE_TITLE)
         app2PackageItemInfo = packageItemInfoWithTitle(APP_2_PACKAGE_NAME, APP_2_PACKAGE_TITLE)
@@ -123,7 +123,7 @@
         val widgetPickerData =
             WidgetPickerData(
                 allWidgets = appTwoWidgetsListBaseEntries(),
-                defaultWidgets = appTwoWidgetsListBaseEntries()
+                defaultWidgets = appTwoWidgetsListBaseEntries(),
             )
 
         val newWidgetData =
@@ -143,19 +143,19 @@
                         addAll(appOneWidgetsListBaseEntries())
                         addAll(appTwoWidgetsListBaseEntries())
                     },
-                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() }
+                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() },
             )
         val recommendations: List<ItemInfo> =
             listOf(
                 PendingAddWidgetInfo(
                     app1WidgetItem1.widgetInfo,
                     CONTAINER_WIDGETS_PREDICTION,
-                    CATEGORY_1
+                    CATEGORY_1,
                 ),
                 PendingAddWidgetInfo(
                     app2WidgetItem1.widgetInfo,
                     CONTAINER_WIDGETS_PREDICTION,
-                    CATEGORY_2
+                    CATEGORY_2,
                 ),
             )
 
@@ -175,7 +175,7 @@
                         addAll(appOneWidgetsListBaseEntries())
                         addAll(appTwoWidgetsListBaseEntries())
                     },
-                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() }
+                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() },
             )
         val recommendations: List<ItemInfo> =
             listOf(
@@ -201,7 +201,7 @@
                         addAll(appTwoWidgetsListBaseEntries())
                     },
                 defaultWidgets = buildList { appTwoWidgetsListBaseEntries() },
-                recommendations = mapOf(CATEGORY_1 to listOf(app1WidgetItem1))
+                recommendations = mapOf(CATEGORY_1 to listOf(app1WidgetItem1)),
             )
 
         val updatedData = widgetPickerData.withRecommendedWidgets(listOf())
@@ -242,7 +242,7 @@
                         addAll(appOneWidgetsListBaseEntries())
                         addAll(appTwoWidgetsListBaseEntries())
                     },
-                defaultWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) }
+                defaultWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) },
             )
         val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
 
@@ -263,7 +263,7 @@
                         addAll(appTwoWidgetsListBaseEntries())
                     },
                 defaultWidgets =
-                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) }
+                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) },
             )
         val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
 
@@ -271,7 +271,7 @@
             findContentEntryForPackageUser(
                 widgetPickerData = widgetPickerData,
                 packageUserKey = app1PackageUserKey,
-                fromDefaultWidgets = true
+                fromDefaultWidgets = true,
             )
 
         assertThat(contentEntry).isNotNull()
@@ -302,7 +302,7 @@
                         addAll(appTwoWidgetsListBaseEntries())
                     },
                 defaultWidgets =
-                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) }
+                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) },
             )
 
         val widgets = findAllWidgetsForPackageUser(widgetPickerData, app1PackageUserKey)
@@ -314,9 +314,7 @@
     @Test
     fun findAllWidgetsForPackageUser_noMatch_returnsEmptyList() {
         val widgetPickerData =
-            WidgetPickerData(
-                allWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) },
-            )
+            WidgetPickerData(allWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) })
         val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
 
         val widgets = findAllWidgetsForPackageUser(widgetPickerData, app1PackageUserKey)
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
index 24d66a3..59f352b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
@@ -41,8 +41,8 @@
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.search.SearchCallback;
@@ -87,7 +87,7 @@
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         doAnswer(invocation -> {
-            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            CachedObject componentWithLabel = invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
         mTestProfile = new InvariantDeviceProfile();
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
index 7adb2b1..2f5fcfe 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
@@ -28,7 +28,6 @@
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.graphics.Point;
 import android.graphics.drawable.Drawable;
 import android.os.UserHandle;
@@ -39,8 +38,9 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.BaseIconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
@@ -99,7 +99,7 @@
         initTestWidgets();
         initTestShortcuts();
 
-        doAnswer(invocation -> ((ComponentWithLabel) invocation.getArgument(0))
+        doAnswer(invocation -> ((CachedObject) invocation.getArgument(0))
                 .getComponent().getPackageName())
                 .when(mIconCache).getTitleNoCache(any());
     }
@@ -280,16 +280,15 @@
     }
 
     private void initTestShortcuts() {
-        PackageManager packageManager = mContext.getPackageManager();
         mShortcut1 = new WidgetItem(new TestShortcutConfigActivityInfo(
                 ComponentName.createRelative(TEST_PACKAGE, ".shortcut1"), UserHandle.CURRENT),
-                mIconCache, packageManager);
+                mIconCache);
         mShortcut2 = new WidgetItem(new TestShortcutConfigActivityInfo(
                 ComponentName.createRelative(TEST_PACKAGE, ".shortcut2"), UserHandle.CURRENT),
-                mIconCache, packageManager);
+                mIconCache);
         mShortcut3 = new WidgetItem(new TestShortcutConfigActivityInfo(
                 ComponentName.createRelative(TEST_PACKAGE, ".shortcut3"), UserHandle.CURRENT),
-                mIconCache, packageManager);
+                mIconCache);
 
     }
 
@@ -300,12 +299,12 @@
         }
 
         @Override
-        public Drawable getFullResIcon(IconCache cache) {
+        public Drawable getFullResIcon(BaseIconCache cache) {
             return null;
         }
 
         @Override
-        public CharSequence getLabel(PackageManager pm) {
+        public CharSequence getLabel() {
             return null;
         }
     }
diff --git a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
index 0dd13a9..b17cd4d 100644
--- a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
@@ -33,17 +33,15 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
-import org.mockito.ArgumentMatchers.any
-import org.mockito.ArgumentMatchers.anyInt
-import org.mockito.ArgumentMatchers.anyList
-import org.mockito.ArgumentMatchers.anyMap
 import org.mockito.Mock
 import org.mockito.Mockito
+import org.mockito.Mockito.mock
 import org.mockito.Mockito.times
 import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations
 import org.mockito.MockitoSession
 import org.mockito.Spy
+import org.mockito.kotlin.any
 import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.spy
@@ -67,7 +65,7 @@
             installedHotseatItems = mutableSetOf("installedHotseatItem"),
             installedWorkspaceItems = mutableSetOf("installedWorkspaceItem"),
             firstScreenInstalledWidgets = mutableSetOf("installedFirstScreenWidget"),
-            secondaryScreenInstalledWidgets = mutableSetOf("installedSecondaryScreenWidget")
+            secondaryScreenInstalledWidgets = mutableSetOf("installedSecondaryScreenWidget"),
         )
     private lateinit var mockitoSession: MockitoSession
 
@@ -75,7 +73,7 @@
     @Mock private lateinit var bgAllAppsList: AllAppsList
     @Mock private lateinit var modelDelegate: ModelDelegate
     @Mock private lateinit var launcherBinder: BaseLauncherBinder
-    @Mock private lateinit var launcherModel: LauncherModel
+    private lateinit var launcherModel: LauncherModel
     @Mock private lateinit var transaction: LoaderTransaction
     @Mock private lateinit var iconCache: IconCache
     @Mock private lateinit var idleLock: LooperIdleLock
@@ -89,6 +87,7 @@
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
+        launcherModel = mock(LauncherModel::class.java)
         mockitoSession =
             ExtendedMockito.mockitoSession()
                 .strictness(Strictness.LENIENT)
@@ -105,15 +104,16 @@
 
         doReturn(TestViewHelpers.findWidgetProvider(false))
             .`when`(context.spyService(AppWidgetManager::class.java))
-            .getAppWidgetInfo(anyInt())
+            .getAppWidgetInfo(any())
         `when`(app.context).thenReturn(context)
         `when`(app.model).thenReturn(launcherModel)
-        `when`(launcherModel.beginLoader(any(LoaderTask::class.java))).thenReturn(transaction)
+
+        `when`(launcherModel.beginLoader(any())).thenReturn(transaction)
         `when`(app.iconCache).thenReturn(iconCache)
         `when`(launcherModel.modelDbController)
             .thenReturn(FactitiousDbController(context, INSERTION_STATEMENT_FILE))
         `when`(app.invariantDeviceProfile).thenReturn(idp)
-        `when`(launcherBinder.newIdleLock(any(LoaderTask::class.java))).thenReturn(idleLock)
+        `when`(launcherBinder.newIdleLock(any())).thenReturn(idleLock)
         `when`(idleLock.awaitLocked(1000)).thenReturn(false)
         `when`(iconCache.updateHandler).thenReturn(iconCacheUpdateHandler)
         context.putObject(UserCache.INSTANCE, userCache)
@@ -149,12 +149,12 @@
 
         verify(launcherBinder).bindWorkspace(true, false)
         verify(modelDelegate).workspaceLoadComplete()
-        verify(modelDelegate).loadAndBindAllAppsItems(any(), any(), any())
+        verify(modelDelegate).loadAndBindAllAppsItems(any(), anyOrNull(), any())
         verify(launcherBinder).bindAllApps()
         verify(iconCacheUpdateHandler, times(4)).updateIcons(any(), any<CachingLogic<Any>>(), any())
         verify(launcherBinder).bindDeepShortcuts()
         verify(launcherBinder).bindWidgets()
-        verify(modelDelegate).loadAndBindOtherItems(any())
+        verify(modelDelegate).loadAndBindOtherItems(anyOrNull())
         verify(iconCacheUpdateHandler).finish()
         verify(modelDelegate).modelLoadComplete()
         verify(transaction).commit()
@@ -209,10 +209,10 @@
         `when`(app.context).thenReturn(spyContext)
         whenever(
                 FirstScreenBroadcastHelper.createModelsForFirstScreenBroadcast(
-                    anyOrNull(),
-                    anyList(),
-                    anyMap(),
-                    anyList()
+                    any(),
+                    any(),
+                    any(),
+                    any(),
                 )
             )
             .thenReturn(listOf(expectedBroadcastModel))
@@ -220,7 +220,7 @@
         whenever(
                 FirstScreenBroadcastHelper.sendBroadcastsForModels(
                     spyContext,
-                    listOf(expectedBroadcastModel)
+                    listOf(expectedBroadcastModel),
                 )
             )
             .thenCallRealMethod()
@@ -239,34 +239,34 @@
         assertEquals(expectedBroadcastModel.installerPackage, actualBroadcastIntent.`package`)
         assertEquals(
             ArrayList(expectedBroadcastModel.installedWorkspaceItems),
-            actualBroadcastIntent.getStringArrayListExtra("workspaceInstalledItems")
+            actualBroadcastIntent.getStringArrayListExtra("workspaceInstalledItems"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.installedHotseatItems),
-            actualBroadcastIntent.getStringArrayListExtra("hotseatInstalledItems")
+            actualBroadcastIntent.getStringArrayListExtra("hotseatInstalledItems"),
         )
         assertEquals(
             ArrayList(
                 expectedBroadcastModel.firstScreenInstalledWidgets +
                     expectedBroadcastModel.secondaryScreenInstalledWidgets
             ),
-            actualBroadcastIntent.getStringArrayListExtra("widgetInstalledItems")
+            actualBroadcastIntent.getStringArrayListExtra("widgetInstalledItems"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.pendingCollectionItems),
-            actualBroadcastIntent.getStringArrayListExtra("folderItem")
+            actualBroadcastIntent.getStringArrayListExtra("folderItem"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.pendingWorkspaceItems),
-            actualBroadcastIntent.getStringArrayListExtra("workspaceItem")
+            actualBroadcastIntent.getStringArrayListExtra("workspaceItem"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.pendingHotseatItems),
-            actualBroadcastIntent.getStringArrayListExtra("hotseatItem")
+            actualBroadcastIntent.getStringArrayListExtra("hotseatItem"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.pendingWidgetItems),
-            actualBroadcastIntent.getStringArrayListExtra("widgetItem")
+            actualBroadcastIntent.getStringArrayListExtra("widgetItem"),
         )
     }
 
@@ -277,10 +277,10 @@
         `when`(app.context).thenReturn(spyContext)
         whenever(
                 FirstScreenBroadcastHelper.createModelsForFirstScreenBroadcast(
-                    anyOrNull(),
-                    anyList(),
-                    anyMap(),
-                    anyList()
+                    any(),
+                    any(),
+                    any(),
+                    any(),
                 )
             )
             .thenReturn(listOf(expectedBroadcastModel))
@@ -288,7 +288,7 @@
         whenever(
                 FirstScreenBroadcastHelper.sendBroadcastsForModels(
                     spyContext,
-                    listOf(expectedBroadcastModel)
+                    listOf(expectedBroadcastModel),
                 )
             )
             .thenCallRealMethod()
@@ -300,7 +300,7 @@
             .runSyncOnBackgroundThread()
 
         // Then
-        verify(spyContext, times(0)).sendBroadcast(any(Intent::class.java))
+        verify(spyContext, times(0)).sendBroadcast(any())
     }
 
     @Test
@@ -310,10 +310,10 @@
         `when`(app.context).thenReturn(spyContext)
         whenever(
                 FirstScreenBroadcastHelper.createModelsForFirstScreenBroadcast(
-                    anyOrNull(),
-                    anyList(),
-                    anyMap(),
-                    anyList()
+                    any(),
+                    any(),
+                    any(),
+                    any(),
                 )
             )
             .thenReturn(listOf(expectedBroadcastModel))
@@ -321,7 +321,7 @@
         whenever(
                 FirstScreenBroadcastHelper.sendBroadcastsForModels(
                     spyContext,
-                    listOf(expectedBroadcastModel)
+                    listOf(expectedBroadcastModel),
                 )
             )
             .thenCallRealMethod()
@@ -334,7 +334,7 @@
             .runSyncOnBackgroundThread()
 
         // Then
-        verify(spyContext, times(0)).sendBroadcast(any(Intent::class.java))
+        verify(spyContext, times(0)).sendBroadcast(any())
     }
 }
 
diff --git a/tests/src/com/android/launcher3/model/PackageUpdatedTaskTest.kt b/tests/src/com/android/launcher3/model/PackageUpdatedTaskTest.kt
index 05f626d..d9af07a 100644
--- a/tests/src/com/android/launcher3/model/PackageUpdatedTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/PackageUpdatedTaskTest.kt
@@ -58,8 +58,7 @@
 @RunWith(AndroidJUnit4::class)
 class PackageUpdatedTaskTest {
 
-    @get:Rule(order = 0) val setFlagsRule = SetFlagsRule()
-    @get:Rule(order = 1) val modelTestRule = ModelTestRule()
+    @get:Rule val setFlagsRule = SetFlagsRule()
 
     private val mUser = UserHandle(0)
     private val mDataModel: BgDataModel = BgDataModel()
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
index e6e02b4..961e7fc 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
@@ -93,7 +93,7 @@
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
         mLauncher.getWorkspace()
                 .openAllWidgets()
-                .getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .getWidget(mWidgetInfo.getLabel())
                 .dragToWorkspace(true, false);
         // Widget id for which the config activity was opened
         mWidgetId = monitor.getWidgetId();
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
index 9c916fa..9a2147a 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
@@ -61,7 +61,7 @@
         WidgetResizeFrame resizeFrame = mLauncher
                 .getWorkspace()
                 .openAllWidgets()
-                .getWidget(widgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .getWidget(widgetInfo.getLabel())
                 .dragWidgetToWorkspace();
 
         assertNotNull("Widget resize frame not shown after widget add", resizeFrame);
@@ -111,7 +111,7 @@
         WidgetResizeFrame resizeFrame = mLauncher
                 .getWorkspace()
                 .openAllWidgets()
-                .getWidget(widgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .getWidget(widgetInfo.getLabel())
                 .dragWidgetToWorkspace();
 
         assertNotNull("Widget resize frame not shown after widget add", resizeFrame);
diff --git a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
index a148744..d653317 100644
--- a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
@@ -16,8 +16,6 @@
 package com.android.launcher3.ui.workspace;
 
 import static com.android.launcher3.util.TestConstants.AppNames.TEST_APP_NAME;
-import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
-import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -40,8 +38,6 @@
 import com.android.launcher3.tapl.HomeAppIconMenuItem;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.Executors;
-import com.android.launcher3.util.rule.ScreenRecordRule;
-import com.android.launcher3.util.rule.TestStabilityRule;
 
 import org.junit.Test;
 
@@ -115,8 +111,6 @@
     }
 
     @Test
-    @TestStabilityRule.Stability(flavors = LOCAL | PLATFORM_POSTSUBMIT) // b/350557998
-    @ScreenRecordRule.ScreenRecord // b/350557998
     public void testShortcutIconWithTheme() throws Exception {
         setThemeEnabled(true);
         initialize(this);
diff --git a/tests/src/com/android/launcher3/util/RoboApiWrapper.kt b/tests/src/com/android/launcher3/util/RoboApiWrapper.kt
index 583652d..7f74e56 100644
--- a/tests/src/com/android/launcher3/util/RoboApiWrapper.kt
+++ b/tests/src/com/android/launcher3/util/RoboApiWrapper.kt
@@ -24,12 +24,10 @@
 
 object RoboApiWrapper {
 
-    fun initialize() {}
-
     fun registerInputStream(
         contentResolver: ContentResolver,
         uri: Uri,
-        inputStreamSupplier: Supplier<InputStream>
+        inputStreamSupplier: Supplier<InputStream>,
     ) {}
 
     fun waitForLooperSync(looper: Looper) {}
diff --git a/tests/src_deviceless/com/android/launcher3/util/RoboApiWrapper.kt b/tests/src_deviceless/com/android/launcher3/util/RoboApiWrapper.kt
index 9232268..a2b8303 100644
--- a/tests/src_deviceless/com/android/launcher3/util/RoboApiWrapper.kt
+++ b/tests/src_deviceless/com/android/launcher3/util/RoboApiWrapper.kt
@@ -16,70 +16,19 @@
 
 package com.android.launcher3.util
 
-import android.content.ComponentName
 import android.content.ContentResolver
-import android.content.Intent
-import android.content.IntentFilter
-import android.content.pm.ApplicationInfo
-import android.content.pm.LauncherActivityInfo
-import android.content.pm.LauncherApps
 import android.net.Uri
 import android.os.Looper
-import android.os.Process
-import androidx.test.platform.app.InstrumentationRegistry
 import java.io.InputStream
 import java.util.function.Supplier
-import org.mockito.Mockito
-import org.mockito.kotlin.whenever
-import org.robolectric.RuntimeEnvironment
 import org.robolectric.Shadows
 
 object RoboApiWrapper {
 
-    fun initialize() {
-        Shadows.shadowOf(
-                RuntimeEnvironment.getApplication().getSystemService(LauncherApps::class.java)
-            )
-            .addEnabledPackage(
-                Process.myUserHandle(),
-                InstrumentationRegistry.getInstrumentation().context.packageName
-            )
-        LauncherModelHelper.ACTIVITY_LIST.forEach {
-            installApp(ComponentName(InstrumentationRegistry.getInstrumentation().context, it))
-        }
-    }
-
-    private fun installApp(componentName: ComponentName) {
-        val app = RuntimeEnvironment.getApplication()
-        val user = Process.myUserHandle()
-
-        val pm = Shadows.shadowOf(app.packageManager)
-        val ai = pm.addActivityIfNotPresent(componentName)
-        pm.addIntentFilterForActivity(
-            componentName,
-            IntentFilter(Intent.ACTION_MAIN).apply { addCategory(Intent.CATEGORY_LAUNCHER) }
-        )
-
-        val li = Mockito.mock(LauncherActivityInfo::class.java)
-        val appInfo = ApplicationInfo().apply { flags = 0 }
-        Mockito.doReturn(ai).whenever(li).activityInfo
-        Mockito.doReturn(appInfo).whenever(li).applicationInfo
-        Mockito.doReturn(user).whenever(li).user
-        Mockito.doReturn(1f).whenever(li).loadingProgress
-        Mockito.doReturn(componentName).whenever(li).componentName
-
-        Shadows.shadowOf(app.getSystemService(LauncherApps::class.java)).apply {
-            addActivity(user, li)
-            addEnabledPackage(user, componentName.packageName)
-            setActivityEnabled(user, componentName)
-            addApplicationInfo(user, componentName.packageName, ai.applicationInfo)
-        }
-    }
-
     fun registerInputStream(
         contentResolver: ContentResolver,
         uri: Uri,
-        inputStreamSupplier: Supplier<InputStream>
+        inputStreamSupplier: Supplier<InputStream>,
     ) {
         Shadows.shadowOf(contentResolver).registerInputStreamSupplier(uri, inputStreamSupplier)
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 6387b05..3097d9c 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -116,8 +116,8 @@
     }
 
     /** Get widget with supplied text. */
-    public Widget getWidget(String labelText) {
-        return getWidget(labelText, null);
+    public Widget getWidget(CharSequence labelText) {
+        return getWidget(labelText.toString(), null);
     }
 
     /** Get widget with supplied text and app package */