diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index c3c42fa..457fdd8 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."
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/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/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/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/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 14c2cc4..7786353 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..5f02893 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;
@@ -1492,6 +1492,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..e15fa54 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,18 +69,25 @@
 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]
  */
 class RecentsWindowManager(context: Context) :
-    RecentsWindowContext(context), RecentsViewContainer, StatefulContainer<RecentsState> {
+    RecentsWindowContext(context),
+    RecentsViewContainer,
+    StatefulContainer<RecentsState>,
+    RecentsAnimationListener {
 
     companion object {
         private const val HOME_APPEAR_DURATION: Long = 250
@@ -98,22 +108,37 @@
     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()
+                }
+            }
+        }
+
     init {
         FallbackWindowInterface.init(this)
+        TaskStackChangeListeners.getInstance().registerTaskStackListener(taskStackChangeListener)
     }
 
     override fun destroy() {
         super.destroy()
+        cleanupRecentsWindow()
         FallbackWindowInterface.getInstance()?.destroy()
+        TaskStackChangeListeners.getInstance().unregisterTaskStackListener(taskStackChangeListener)
+        callbacks?.removeListener(this)
     }
 
     override fun startHome() {
@@ -135,6 +160,7 @@
                 ),
             )
         OverviewComponentObserver.startHomeIntentSafely(this, options.toBundle(), TAG)
+        stateManager.moveToRestState()
     }
 
     private val mAnimationToHomeFactory =
@@ -164,29 +190,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(this)
     }
 
-    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 +247,25 @@
 
         mSystemUiController = SystemUiController(windowView)
         onInitListener?.test(true)
+
+        this.callbacks = callbacks
+        callbacks?.addListener(this)
+    }
+
+    override fun onRecentsAnimationCanceled(thumbnailDatas: HashMap<Int, ThumbnailData>) {
+        super.onRecentsAnimationCanceled(thumbnailDatas)
+        recentAnimationStopped()
+    }
+
+    override fun onRecentsAnimationFinished(controller: RecentsAnimationController) {
+        super.onRecentsAnimationFinished(controller)
+        recentAnimationStopped()
+    }
+
+    private fun recentAnimationStopped() {
+        if (isInState(BACKGROUND_APP)) {
+            cleanupRecentsWindow()
+        }
     }
 
     override fun canStartHomeSafely(): Boolean {
@@ -246,14 +297,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 +374,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/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/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/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..741be50 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
@@ -16,46 +16,19 @@
 
 package com.android.launcher3.taskbar.rules
 
-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 com.android.launcher3.util.MainThreadInitializedObject.ObjectSandbox
+import com.android.launcher3.util.SandboxApplication
+import org.junit.rules.TestRule
 
-/**
- * Sandbox wrapper where [createWindowContext] provides contexts that are still sandboxed within
- * [application].
- *
- * 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].
- */
-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
+/** Sandbox Context for running Taskbar tests. */
+class TaskbarWindowSandboxContext private constructor(base: SandboxApplication) :
+    ContextWrapper(base), ObjectSandbox by base, TestRule by base {
 
     companion object {
-        /** Creates a [TaskbarWindowSandboxContext] to sandbox [base] for Taskbar tests. */
-        fun create(base: Context): TaskbarWindowSandboxContext {
-            return SandboxContext(base).let { TaskbarWindowSandboxContext(it, it) }
+        /** Creates a [SandboxApplication] for Taskbar tests. */
+        fun create(): TaskbarWindowSandboxContext {
+            return TaskbarWindowSandboxContext(SandboxApplication())
         }
     }
 }
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
deleted file mode 100644
index 4834d48..0000000
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
+++ /dev/null
@@ -1,45 +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.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.google.common.truth.Truth.assertThat
-import org.junit.Test
-import org.junit.runner.RunWith
-
-@RunWith(LauncherMultivalentJUnit::class)
-@LauncherMultivalentJUnit.EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
-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)
-    }
-
-    @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)
-    }
-}
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/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/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index b0ec9b0..200d5a7 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -1872,13 +1872,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,
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..e7b9d89
--- /dev/null
+++ b/src/com/android/launcher3/LauncherModel.kt
@@ -0,0 +1,601 @@
+/*
+ * 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.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.WorkerThread
+import com.android.launcher3.celllayout.CellPosMapper
+import com.android.launcher3.config.FeatureFlags
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.icons.cache.BaseIconCache
+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.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.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.testing.shared.TestProtocol.sDebugTracing
+import com.android.launcher3.util.ApplicationInfoWrapper
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import com.android.launcher3.util.Executors.MODEL_EXECUTOR
+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.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.
+ */
+class LauncherModel(
+    private val context: Context,
+    private val mApp: LauncherAppState,
+    private val iconCache: IconCache,
+    private val appFilter: AppFilter,
+    private val mPmHelper: PackageManagerHelper,
+    isPrimaryInstance: Boolean,
+) : InstallSessionTracker.Callback {
+
+    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))
+        val pinnedShortcuts: List<ShortcutInfo> =
+            ShortcutRequest(context, user).forPackage(packageName).query(ShortcutRequest.PINNED)
+        if (pinnedShortcuts.isNotEmpty()) {
+            enqueueModelUpdateTask(ShortcutsChangedTask(packageName, pinnedShortcuts, 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")
+        val action = intent.action
+        if (Intent.ACTION_LOCALE_CHANGED == action) {
+            // If we have changed locale we need to clear out the labels in all apps/workspace.
+            forceReload()
+        } else if (DevicePolicyManager.ACTION_DEVICE_POLICY_RESOURCE_UPDATED == action) {
+            enqueueModelUpdateTask(ReloadStringCacheTask(this.modelDelegate))
+        } else if (BuildConfig.IS_STUDIO_BUILD && LauncherAppState.ACTION_FORCE_ROLOAD == action) {
+            forceReload()
+        }
+    }
+
+    /**
+     * Called then there use a user event
+     *
+     * @see UserCache.addUserEventListener
+     */
+    fun onUserEvent(user: UserHandle, action: String) {
+        if (Intent.ACTION_MANAGED_PROFILE_AVAILABLE == action && mShouldReloadWorkProfile) {
+            mShouldReloadWorkProfile = false
+            forceReload()
+        } else if (
+            Intent.ACTION_MANAGED_PROFILE_AVAILABLE == action ||
+                Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE == action
+        ) {
+            mShouldReloadWorkProfile = false
+            enqueueModelUpdateTask(
+                PackageUpdatedTask(PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user)
+            )
+        } else if (
+            UserCache.ACTION_PROFILE_LOCKED == action || UserCache.ACTION_PROFILE_UNLOCKED == action
+        ) {
+            enqueueModelUpdateTask(
+                UserLockStateChangedTask(user, UserCache.ACTION_PROFILE_UNLOCKED == action)
+            )
+        } else if (
+            UserCache.ACTION_PROFILE_ADDED == action || UserCache.ACTION_PROFILE_REMOVED == action
+        ) {
+            forceReload()
+        } else if (
+            UserCache.ACTION_PROFILE_AVAILABLE == action ||
+                UserCache.ACTION_PROFILE_UNAVAILABLE == 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(
+                PackageUpdatedTask(PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user)
+            )
+        }
+        if (Intent.ACTION_MANAGED_PROFILE_REMOVED == action) {
+            LauncherPrefs.get(mApp.context).put(LauncherPrefs.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
+     */
+    fun 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 */
+    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 {
+                    stopLoader()
+                    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) }
+    }
+
+    override fun onInstallSessionCreated(sessionInfo: PackageInstallInfo) {
+        if (FeatureFlags.PROMISE_APPS_IN_ALL_APPS.get()) {
+            enqueueModelUpdateTask { taskController, _, apps ->
+                apps.addPromiseApp(mApp.context, sessionInfo)
+                taskController.bindApplicationsIfNeeded()
+            }
+        }
+    }
+
+    override fun onSessionFailure(packageName: String, user: UserHandle) {
+        enqueueModelUpdateTask { taskController, dataModel, apps ->
+            val iconCache = mApp.iconCache
+            val removedIds = IntSet()
+            val archivedWorkspaceItemsToCacheRefresh = HashSet<WorkspaceItemInfo>()
+            val isAppArchived = ApplicationInfoWrapper(mApp.context, packageName, user).isArchived()
+            synchronized(dataModel) {
+                if (isAppArchived) {
+                    // Remove package icon cache entry for archived app in case of a session
+                    // failure.
+                    mApp.iconCache.remove(
+                        ComponentName(packageName, packageName + BaseIconCache.EMPTY_CLASS_NAME),
+                        user,
+                    )
+                }
+                for (info in dataModel.itemsIdMap) {
+                    if (
+                        (info is WorkspaceItemInfo && info.hasPromiseIconUi()) &&
+                            user == info.user &&
+                            info.intent != null
+                    ) {
+                        if (TextUtils.equals(packageName, info.intent!!.getPackage())) {
+                            removedIds.add(info.id)
+                        }
+                        if (info.isArchived()) {
+                            // Refresh icons on the workspace for archived apps.
+                            iconCache.getTitleAndIcon(info, info.usingLowResIcon())
+                            archivedWorkspaceItemsToCacheRefresh.add(info)
+                        }
+                    }
+                }
+                if (isAppArchived) {
+                    apps.updateIconsAndLabels(hashSetOf(packageName), user)
+                }
+            }
+
+            if (!removedIds.isEmpty && !isAppArchived) {
+                taskController.deleteAndBindComponentsRemoved(
+                    ItemInfoMatcher.ofItemIds(removedIds),
+                    "removed because install session failed",
+                )
+            }
+            if (archivedWorkspaceItemsToCacheRefresh.isNotEmpty()) {
+                taskController.bindUpdatedWorkspaceItems(
+                    archivedWorkspaceItemsToCacheRefresh.stream().toList()
+                )
+            }
+            if (isAppArchived) {
+                taskController.bindApplicationsIfNeeded()
+            }
+        }
+    }
+
+    override fun onPackageStateChanged(installInfo: PackageInstallInfo) {
+        enqueueModelUpdateTask(PackageInstallStateChangedTask(installInfo))
+    }
+
+    /** Updates the icons and label of all pending icons for the provided package name. */
+    override fun onUpdateSessionDisplay(key: PackageUserKey, info: PackageInstaller.SessionInfo) {
+        mApp.iconCache.updateSessionCache(key, info)
+
+        val packages = HashSet<String>()
+        packages.add(key.mPackageName)
+        enqueueModelUpdateTask(
+            CacheDataUpdatedTask(CacheDataUpdatedTask.OP_SESSION_UPDATE, key.mUser, packages)
+        )
+    }
+
+    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) {
+        updateAndBindWorkspaceItem {
+            si.updateFromDeepShortcutInfo(info, mApp.context)
+            mApp.iconCache.getShortcutIcon(si, info)
+            si
+        }
+    }
+
+    /** Utility method to update a shortcut on the background thread. */
+    private fun updateAndBindWorkspaceItem(itemProvider: Supplier<WorkspaceItemInfo>) {
+        enqueueModelUpdateTask { taskController, _, _ ->
+            val info = itemProvider.get()
+            taskController.getModelWriter().updateItemInDatabase(info)
+            taskController.bindUpdatedWorkspaceItems(listOf(info))
+        }
+    }
+
+    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/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/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/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/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/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/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/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/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/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/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/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/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)
     }