Merge "Add permanent TaskView logs to track taskIds being launched" into main
diff --git a/go/quickstep/src/com/android/launcher3/model/AppShareabilityManager.java b/go/quickstep/src/com/android/launcher3/model/AppShareabilityManager.java
index 0d0f700..556d29c 100644
--- a/go/quickstep/src/com/android/launcher3/model/AppShareabilityManager.java
+++ b/go/quickstep/src/com/android/launcher3/model/AppShareabilityManager.java
@@ -35,6 +35,7 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.launcher3.model.AppShareabilityDatabase.ShareabilityDao;
 import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.SafeCloseable;
 
 import java.lang.annotation.Retention;
 import java.util.ArrayList;
@@ -47,7 +48,7 @@
  * Each app's status is retrieved from the Play Store's API. Statuses are cached in order
  * to limit extraneous calls to that API (which can be time-consuming).
  */
-public class AppShareabilityManager {
+public class AppShareabilityManager implements SafeCloseable {
     @Retention(SOURCE)
     @IntDef({
         ShareabilityStatus.UNKNOWN,
@@ -194,6 +195,11 @@
         }
     }
 
+    @Override
+    public void close() {
+        mDatabase.close();
+    }
+
     /**
      * Provides a testable instance of this class
      * This instance allows database queries on the main thread
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index fb14f9e..65a49bd 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -326,8 +326,12 @@
         super.destroy();
         mActive = false;
         StatsLogCompatManager.LOGS_CONSUMER.remove(mAppEventProducer);
-        if (mIsPrimaryInstance) {
-            mStatsManager.clearPullAtomCallback(SysUiStatsLog.LAUNCHER_LAYOUT_SNAPSHOT);
+        if (mIsPrimaryInstance && mStatsManager != null) {
+            try {
+                mStatsManager.clearPullAtomCallback(SysUiStatsLog.LAUNCHER_LAYOUT_SNAPSHOT);
+            } catch (RuntimeException e) {
+                Log.e(TAG, "Failed to unregister snapshot logging callback with StatsManager", e);
+            }
         }
         destroyPredictors();
     }
diff --git a/quickstep/src/com/android/launcher3/model/WellbeingModel.java b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
index c345d6e..a7c9652 100644
--- a/quickstep/src/com/android/launcher3/model/WellbeingModel.java
+++ b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
@@ -32,7 +32,6 @@
 import android.os.Bundle;
 import android.os.DeadObjectException;
 import android.os.Handler;
-import android.os.Looper;
 import android.os.Process;
 import android.os.UserHandle;
 import android.text.TextUtils;
@@ -48,9 +47,10 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.popup.RemoteActionShortcut;
 import com.android.launcher3.popup.SystemShortcut;
-import com.android.launcher3.util.BgObjectWithLooper;
+import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.SimpleBroadcastReceiver;
 import com.android.launcher3.views.ActivityContext;
 
@@ -61,7 +61,7 @@
 /**
  * Data model for digital wellbeing status of apps.
  */
-public final class WellbeingModel extends BgObjectWithLooper {
+public final class WellbeingModel implements SafeCloseable {
     private static final String TAG = "WellbeingModel";
     private static final int[] RETRY_TIMES_MS = {5000, 15000, 30000};
     private static final boolean DEBUG = false;
@@ -81,8 +81,12 @@
     private final Context mContext;
     private final String mWellbeingProviderPkg;
 
-    private Handler mWorkerHandler;
-    private ContentObserver mContentObserver;
+    private final Handler mWorkerHandler;
+    private final ContentObserver mContentObserver;
+    private final SimpleBroadcastReceiver mWellbeingAppChangeReceiver =
+            new SimpleBroadcastReceiver(t -> restartObserver());
+    private final SimpleBroadcastReceiver mAppAddRemoveReceiver =
+            new SimpleBroadcastReceiver(this::onAppPackageChanged);
 
     private final Object mModelLock = new Object();
     // Maps the action Id to the corresponding RemoteAction
@@ -94,16 +98,23 @@
     private WellbeingModel(final Context context) {
         mContext = context;
         mWellbeingProviderPkg = mContext.getString(R.string.wellbeing_provider_pkg);
-        initializeInBackground("WellbeingHandler");
+        mWorkerHandler = new Handler(TextUtils.isEmpty(mWellbeingProviderPkg)
+                ? Executors.UI_HELPER_EXECUTOR.getLooper()
+                : Executors.getPackageExecutor(mWellbeingProviderPkg).getLooper());
+
+        mContentObserver = new ContentObserver(mWorkerHandler) {
+            @Override
+            public void onChange(boolean selfChange, Uri uri) {
+                updateAllPackages();
+            }
+        };
+        mWorkerHandler.post(this::initializeInBackground);
     }
 
-    @Override
-    protected void onInitialized(Looper looper) {
-        mWorkerHandler = new Handler(looper);
-        mContentObserver = newContentObserver(mWorkerHandler, this::onWellbeingUriChanged);
+    private void initializeInBackground() {
         if (!TextUtils.isEmpty(mWellbeingProviderPkg)) {
             mContext.registerReceiver(
-                    new SimpleBroadcastReceiver(t -> restartObserver()),
+                    mWellbeingAppChangeReceiver,
                     getPackageFilter(mWellbeingProviderPkg,
                             Intent.ACTION_PACKAGE_ADDED, Intent.ACTION_PACKAGE_CHANGED,
                             Intent.ACTION_PACKAGE_REMOVED, Intent.ACTION_PACKAGE_DATA_CLEARED,
@@ -113,17 +124,21 @@
             IntentFilter filter = new IntentFilter(Intent.ACTION_PACKAGE_ADDED);
             filter.addAction(Intent.ACTION_PACKAGE_REMOVED);
             filter.addDataScheme("package");
-            mContext.registerReceiver(new SimpleBroadcastReceiver(this::onAppPackageChanged),
-                    filter, null, mWorkerHandler);
+            mContext.registerReceiver(mAppAddRemoveReceiver, filter, null, mWorkerHandler);
 
             restartObserver();
         }
     }
 
-    @WorkerThread
-    private void onWellbeingUriChanged(Uri uri) {
-        Preconditions.assertNonUiThread();
-        updateAllPackages();
+    @Override
+    public void close() {
+        if (!TextUtils.isEmpty(mWellbeingProviderPkg)) {
+            mWorkerHandler.post(() -> {
+                mWellbeingAppChangeReceiver.unregisterReceiverSafely(mContext);
+                mAppAddRemoveReceiver.unregisterReceiverSafely(mContext);
+                mContext.getContentResolver().unregisterContentObserver(mContentObserver);
+            });
+        }
     }
 
     public void setInTest(boolean inTest) {
diff --git a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
index 74376c8..5ac5761 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
@@ -37,6 +37,9 @@
 import com.android.quickstep.views.DesktopAppSelectView;
 import com.android.wm.shell.desktopmode.IDesktopTaskListener;
 
+import java.util.HashSet;
+import java.util.Set;
+
 /**
  * Controls the visibility of the workspace and the resumed / paused state when desktop mode
  * is enabled.
@@ -48,6 +51,7 @@
     private static final boolean IS_STASHING_ENABLED = SystemProperties.getBoolean(
             "persist.wm.debug.desktop_stashing", false);
     private final Launcher mLauncher;
+    private final Set<DesktopVisibilityListener> mDesktopVisibilityListeners = new HashSet<>();
 
     private int mVisibleDesktopTasksCount;
     private boolean mInOverviewState;
@@ -127,6 +131,16 @@
         return mVisibleDesktopTasksCount;
     }
 
+    /** Registers a listener for Desktop Mode visibility updates. */
+    public void registerDesktopVisibilityListener(DesktopVisibilityListener listener) {
+        mDesktopVisibilityListeners.add(listener);
+    }
+
+    /** Removes a previously registered Desktop Mode visibility listener. */
+    public void unregisterDesktopVisibilityListener(DesktopVisibilityListener listener) {
+        mDesktopVisibilityListeners.remove(listener);
+    }
+
     /**
      * Sets the number of desktop windows that are visible and updates launcher visibility based on
      * it.
@@ -140,7 +154,12 @@
         if (visibleTasksCount != mVisibleDesktopTasksCount) {
             final boolean wasVisible = mVisibleDesktopTasksCount > 0;
             final boolean isVisible = visibleTasksCount > 0;
+            final boolean wereDesktopTasksVisibleBefore = areDesktopTasksVisible();
             mVisibleDesktopTasksCount = visibleTasksCount;
+            final boolean areDesktopTasksVisibleNow = areDesktopTasksVisible();
+            if (wereDesktopTasksVisibleBefore != areDesktopTasksVisibleNow) {
+                notifyDesktopVisibilityListeners(areDesktopTasksVisibleNow);
+            }
 
             if (!enableDesktopWindowingWallpaperActivity() && wasVisible != isVisible) {
                 // TODO: b/333533253 - Remove after flag rollout
@@ -179,15 +198,22 @@
                     + " currentValue=" + mInOverviewState);
         }
         if (overviewStateEnabled != mInOverviewState) {
+            final boolean wereDesktopTasksVisibleBefore = areDesktopTasksVisible();
             mInOverviewState = overviewStateEnabled;
+            final boolean areDesktopTasksVisibleNow = areDesktopTasksVisible();
+            if (wereDesktopTasksVisibleBefore != areDesktopTasksVisibleNow) {
+                notifyDesktopVisibilityListeners(areDesktopTasksVisibleNow);
+            }
+
             if (enableDesktopWindowingWallpaperActivity()) {
                 return;
             }
             // TODO: b/333533253 - Clean up after flag rollout
+
             if (mInOverviewState) {
                 setLauncherViewsVisibility(View.VISIBLE);
                 markLauncherResumed();
-            } else if (areDesktopTasksVisible() && !mGestureInProgress) {
+            } else if (areDesktopTasksVisibleNow && !mGestureInProgress) {
                 // Switching out of overview state and gesture finished.
                 // If desktop tasks are still visible, hide launcher again.
                 setLauncherViewsVisibility(View.INVISIBLE);
@@ -196,6 +222,15 @@
         }
     }
 
+    private void notifyDesktopVisibilityListeners(boolean areDesktopTasksVisible) {
+        if (DEBUG) {
+            Log.d(TAG, "notifyDesktopVisibilityListeners: visible=" + areDesktopTasksVisible);
+        }
+        for (DesktopVisibilityListener listener : mDesktopVisibilityListeners) {
+            listener.onDesktopVisibilityChanged(areDesktopTasksVisible);
+        }
+    }
+
     /**
      * TODO: b/333533253 - Remove after flag rollout
      */
@@ -359,4 +394,14 @@
         mSelectAppToast.hide();
         mSelectAppToast = null;
     }
+
+    /** A listener for when the user enters/exits Desktop Mode. */
+    public interface DesktopVisibilityListener {
+        /**
+         * Callback for when the user enters or exits Desktop Mode
+         *
+         * @param visible whether Desktop Mode is now visible
+         */
+        void onDesktopVisibilityChanged(boolean visible);
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/DesktopTaskbarRunningAppsController.kt b/quickstep/src/com/android/launcher3/taskbar/DesktopTaskbarRunningAppsController.kt
new file mode 100644
index 0000000..f665e21
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/DesktopTaskbarRunningAppsController.kt
@@ -0,0 +1,139 @@
+/*
+ * 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
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration
+import android.util.Log
+import android.util.SparseArray
+import androidx.annotation.VisibleForTesting
+import androidx.core.util.valueIterator
+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.statehandlers.DesktopVisibilityController
+import com.android.quickstep.RecentsModel
+import kotlin.collections.filterNotNull
+
+/**
+ * Shows running apps when in Desktop Mode.
+ *
+ * Users can enter and exit Desktop Mode at run-time, meaning this class falls back to the default
+ * recent-apps behaviour when outside of Desktop Mode.
+ *
+ * This class should only be used if
+ * [com.android.window.flags.Flags.enableDesktopWindowingTaskbarRunningApps] is enabled.
+ */
+class DesktopTaskbarRunningAppsController(
+    private val recentsModel: RecentsModel,
+    private val desktopVisibilityController: DesktopVisibilityController?,
+) : TaskbarRecentAppsController() {
+
+    private var apps: Array<AppInfo>? = null
+    private var allRunningDesktopAppInfos: List<AppInfo>? = null
+    private var runningDesktopAppInfosExceptHotseatItems: List<ItemInfo>? = null
+
+    private val isInDesktopMode: Boolean
+        get() = desktopVisibilityController?.areDesktopTasksVisible() ?: false
+
+    override fun onDestroy() {
+        super.onDestroy()
+        apps = null
+    }
+
+    @VisibleForTesting
+    public override fun setApps(apps: Array<AppInfo>?) {
+        this.apps = apps
+    }
+
+    override fun isEnabled() = true
+
+    @VisibleForTesting
+    public override fun updateHotseatItemInfos(hotseatItems: Array<ItemInfo>?): Array<ItemInfo>? {
+        val actualHotseatItems = hotseatItems ?: return super.updateHotseatItemInfos(null)
+        if (!isInDesktopMode) {
+            Log.d(TAG, "updateHotseatItemInfos: not in Desktop Mode")
+            return hotseatItems
+        }
+        val newHotseatItemInfos =
+            actualHotseatItems
+                // Ignore predicted apps - we show running apps instead
+                .filter { itemInfo -> !itemInfo.isPredictedItem }
+                .toMutableList()
+        val runningDesktopAppInfos =
+            runningDesktopAppInfosExceptHotseatItems ?: return newHotseatItemInfos.toTypedArray()
+        newHotseatItemInfos.addAll(runningDesktopAppInfos)
+        return newHotseatItemInfos.toTypedArray()
+    }
+
+    @VisibleForTesting
+    public override fun updateRunningApps(hotseatItems: SparseArray<ItemInfo>?) {
+        if (!isInDesktopMode) {
+            Log.d(TAG, "updateRunningApps: not in Desktop Mode")
+            mControllers.taskbarViewController.commitRunningAppsToUI()
+            return
+        }
+        val allRunningDesktopAppInfos = getRunningDesktopAppInfos()
+        this.allRunningDesktopAppInfos = allRunningDesktopAppInfos
+        runningDesktopAppInfosExceptHotseatItems =
+            hotseatItems?.let {
+                getRunningDesktopAppInfosExceptHotseatApps(allRunningDesktopAppInfos, it.toList())
+            }
+
+        mControllers.taskbarViewController.commitRunningAppsToUI()
+    }
+
+    private fun getRunningDesktopAppInfosExceptHotseatApps(
+        allRunningDesktopAppInfos: List<AppInfo>,
+        hotseatItems: List<ItemInfo>
+    ): List<ItemInfo> {
+        val hotseatPackages = hotseatItems.map { it.targetPackage }
+        return allRunningDesktopAppInfos
+            .filter { appInfo -> !hotseatPackages.contains(appInfo.targetPackage) }
+            .map { WorkspaceItemInfo(it) }
+    }
+
+    private fun getRunningDesktopAppInfos(): List<AppInfo> {
+        return getAppInfosFromRunningTasks(
+            recentsModel.runningTasks
+                .filter { taskInfo: RunningTaskInfo ->
+                    taskInfo.windowingMode == WindowConfiguration.WINDOWING_MODE_FREEFORM
+                }
+                .toList()
+        )
+    }
+
+    // TODO(b/335398876) fetch app icons from Tasks instead of AppInfos
+    private fun getAppInfosFromRunningTasks(tasks: List<RunningTaskInfo>): List<AppInfo> {
+        // Early return if apps is empty, since we then have no AppInfo to compare to
+        if (apps == null) {
+            return emptyList()
+        }
+        val packageNames = tasks.map { it.realActivity?.packageName }.distinct().filterNotNull()
+        return packageNames
+            .map { packageName -> apps?.find { app -> packageName == app.targetPackage } }
+            .filterNotNull()
+    }
+
+    private fun <E> SparseArray<E>.toList(): List<E> {
+        return valueIterator().asSequence().toList()
+    }
+
+    companion object {
+        private const val TAG = "TabletDesktopTaskbarRunningAppsController"
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 49d4afe..679528b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -41,6 +41,8 @@
 import static com.android.quickstep.util.AnimUtils.completeRunnableListCallback;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_VOICE_INTERACTION_WINDOW_SHOWING;
+import static com.android.window.flags.Flags.enableDesktopWindowingMode;
+import static com.android.window.flags.Flags.enableDesktopWindowingTaskbarRunningApps;
 
 import android.animation.AnimatorSet;
 import android.animation.ValueAnimator;
@@ -127,7 +129,9 @@
 import com.android.launcher3.util.VibratorWrapper;
 import com.android.launcher3.util.ViewCache;
 import com.android.launcher3.views.ActivityContext;
+import com.android.quickstep.LauncherActivityInterface;
 import com.android.quickstep.NavHandle;
+import com.android.quickstep.RecentsModel;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
@@ -244,7 +248,7 @@
 
         mAccessibilityDelegate = new TaskbarShortcutMenuAccessibilityDelegate(this);
 
-        final boolean isDesktopMode = getPackageManager().hasSystemFeature(FEATURE_PC);
+        final boolean isPcMode = getPackageManager().hasSystemFeature(FEATURE_PC);
 
         // If Bubble bar is present, TaskbarControllers depends on it so build it first.
         Optional<BubbleControllers> bubbleControllersOptional = Optional.empty();
@@ -276,7 +280,7 @@
         mControllers = new TaskbarControllers(this,
                 new TaskbarDragController(this),
                 buttonController,
-                isDesktopMode
+                isPcMode
                         ? new DesktopNavbarButtonsViewController(this, mNavigationBarPanelContext,
                                 navButtonsView)
                         : new NavbarButtonsViewController(this, mNavigationBarPanelContext,
@@ -301,9 +305,7 @@
                 new VoiceInteractionWindowController(this),
                 new TaskbarTranslationController(this),
                 new TaskbarSpringOnStashController(this),
-                isDesktopMode
-                        ? new DesktopTaskbarRecentAppsController(this)
-                        : TaskbarRecentAppsController.DEFAULT,
+                createTaskbarRecentAppsController(isPcMode),
                 new TaskbarEduTooltipController(this),
                 new KeyboardQuickSwitchController(),
                 new TaskbarPinningController(this),
@@ -312,6 +314,18 @@
         mLauncherPrefs = LauncherPrefs.get(this);
     }
 
+    private TaskbarRecentAppsController createTaskbarRecentAppsController(boolean isPcMode) {
+        if (isPcMode) return new DesktopTaskbarRecentAppsController(this);
+        // TODO(b/335401172): unify DesktopMode checks in Launcher
+        final boolean showRunningAppsInDesktopMode = enableDesktopWindowingMode()
+                && enableDesktopWindowingTaskbarRunningApps();
+        return showRunningAppsInDesktopMode
+                        ? new DesktopTaskbarRunningAppsController(
+                                RecentsModel.INSTANCE.get(this),
+                                LauncherActivityInterface.INSTANCE.getDesktopVisibilityController())
+                        : TaskbarRecentAppsController.DEFAULT;
+    }
+
     /** Updates {@link DeviceProfile} instances for any Taskbar windows. */
     public void updateDeviceProfile(DeviceProfile launcherDp) {
         applyDeviceProfile(launcherDp);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java
index 14d46d1..6c84f80 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java
@@ -15,6 +15,9 @@
  */
 package com.android.launcher3.taskbar;
 
+import static com.android.window.flags.Flags.enableDesktopWindowingMode;
+import static com.android.window.flags.Flags.enableDesktopWindowingTaskbarRunningApps;
+
 import android.util.SparseArray;
 import android.view.View;
 
@@ -26,6 +29,7 @@
 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.statehandlers.DesktopVisibilityController;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
@@ -33,6 +37,7 @@
 import com.android.launcher3.util.LauncherBindableItemsContainer;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
+import com.android.quickstep.LauncherActivityInterface;
 import com.android.quickstep.RecentsModel;
 
 import java.io.PrintWriter;
@@ -62,6 +67,8 @@
     // Used to defer any UI updates during the SUW unstash animation.
     private boolean mDeferUpdatesForSUW;
     private Runnable mDeferredUpdates;
+    private DesktopVisibilityController.DesktopVisibilityListener mDesktopVisibilityListener =
+            visible -> updateRunningApps();
 
     public TaskbarModelCallbacks(
             TaskbarActivityContext context, TaskbarView container) {
@@ -73,6 +80,15 @@
         mControllers = controllers;
         if (mControllers.taskbarRecentAppsController.isEnabled()) {
             RecentsModel.INSTANCE.get(mContext).registerRunningTasksListener(this);
+
+            if (shouldShowRunningAppsInDesktopMode()) {
+                DesktopVisibilityController desktopVisibilityController =
+                        LauncherActivityInterface.INSTANCE.getDesktopVisibilityController();
+                if (desktopVisibilityController != null) {
+                    desktopVisibilityController.registerDesktopVisibilityListener(
+                            mDesktopVisibilityListener);
+                }
+            }
         }
     }
 
@@ -81,6 +97,20 @@
      */
     public void unregisterListeners() {
         RecentsModel.INSTANCE.get(mContext).unregisterRunningTasksListener();
+
+        if (shouldShowRunningAppsInDesktopMode()) {
+            DesktopVisibilityController desktopVisibilityController =
+                    LauncherActivityInterface.INSTANCE.getDesktopVisibilityController();
+            if (desktopVisibilityController != null) {
+                desktopVisibilityController.unregisterDesktopVisibilityListener(
+                        mDesktopVisibilityListener);
+            }
+        }
+    }
+
+    private boolean shouldShowRunningAppsInDesktopMode() {
+        // TODO(b/335401172): unify DesktopMode checks in Launcher
+        return enableDesktopWindowingMode() && enableDesktopWindowingTaskbarRunningApps();
     }
 
     @Override
diff --git a/quickstep/src/com/android/quickstep/RotationTouchHelper.java b/quickstep/src/com/android/quickstep/RotationTouchHelper.java
index bf50d70..3380291 100644
--- a/quickstep/src/com/android/quickstep/RotationTouchHelper.java
+++ b/quickstep/src/com/android/quickstep/RotationTouchHelper.java
@@ -39,6 +39,7 @@
 import com.android.launcher3.util.DisplayController.Info;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.NavigationMode;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.quickstep.util.RecentsOrientedState;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.shared.system.TaskStackChangeListener;
@@ -50,7 +51,7 @@
 /**
  * Helper class for transforming touch events
  */
-public class RotationTouchHelper implements DisplayInfoChangeListener {
+public class RotationTouchHelper implements DisplayInfoChangeListener, SafeCloseable {
 
     public static final MainThreadInitializedObject<RotationTouchHelper> INSTANCE =
             new MainThreadInitializedObject<>(RotationTouchHelper::new);
@@ -197,6 +198,11 @@
         mOnDestroyActions.add(action);
     }
 
+    @Override
+    public void close() {
+        destroy();
+    }
+
     /**
      * Cleans up all the registered listeners and receivers.
      */
diff --git a/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java b/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java
index f474796..29a57fc 100644
--- a/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java
+++ b/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java
@@ -24,22 +24,30 @@
 
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.SafeCloseable;
 
 public class SimpleOrientationTouchTransformer implements
-        DisplayController.DisplayInfoChangeListener {
+        DisplayController.DisplayInfoChangeListener, SafeCloseable {
 
     public static final MainThreadInitializedObject<SimpleOrientationTouchTransformer> INSTANCE =
             new MainThreadInitializedObject<>(SimpleOrientationTouchTransformer::new);
 
+    private final Context mContext;
     private OrientationRectF mOrientationRectF;
 
     public SimpleOrientationTouchTransformer(Context context) {
+        mContext = context;
         DisplayController.INSTANCE.get(context).addChangeListener(this);
         onDisplayInfoChanged(context, DisplayController.INSTANCE.get(context).getInfo(),
                 CHANGE_ALL);
     }
 
     @Override
+    public void close() {
+        DisplayController.INSTANCE.get(mContext).removeChangeListener(this);
+    }
+
+    @Override
     public void onDisplayInfoChanged(Context context, DisplayController.Info info, int flags) {
         if ((flags & (CHANGE_ROTATION | CHANGE_ACTIVE_SCREEN)) == 0) {
             return;
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index 30bb863..ab9840f 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -16,6 +16,7 @@
 package com.android.quickstep;
 
 import static android.app.ActivityManager.RECENT_IGNORE_UNAVAILABLE;
+import static android.content.pm.PackageManager.FEATURE_PC;
 
 import static com.android.launcher3.Flags.enableUnfoldStateAnimation;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
@@ -23,6 +24,8 @@
 import static com.android.launcher3.util.SplitConfigurationOptions.StagePosition;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.RECENT_TASKS_MISSING;
 import static com.android.quickstep.util.LogUtils.splitFailureMessage;
+import static com.android.window.flags.Flags.enableDesktopWindowingMode;
+import static com.android.window.flags.Flags.enableDesktopWindowingTaskbarRunningApps;
 
 import android.app.ActivityManager;
 import android.app.ActivityOptions;
@@ -32,7 +35,6 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ActivityInfo;
-import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
 import android.graphics.Point;
 import android.graphics.Rect;
@@ -65,6 +67,7 @@
 import com.android.internal.view.AppearanceRegion;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.quickstep.util.ActiveGestureLog;
 import com.android.quickstep.util.AssistUtils;
 import com.android.quickstep.util.unfold.ProxyUnfoldTransitionProvider;
@@ -108,7 +111,7 @@
 /**
  * Holds the reference to SystemUI.
  */
-public class SystemUiProxy implements ISystemUiProxy, NavHandle {
+public class SystemUiProxy implements ISystemUiProxy, NavHandle, SafeCloseable {
     private static final String TAG = SystemUiProxy.class.getSimpleName();
 
     public static final MainThreadInitializedObject<SystemUiProxy> INSTANCE =
@@ -199,6 +202,9 @@
     }
 
     @Override
+    public void close() { }
+
+    @Override
     public void onBackPressed() {
         if (mSystemUiProxy != null) {
             try {
@@ -1366,8 +1372,7 @@
      * Gets the set of running tasks.
      */
     public ArrayList<ActivityManager.RunningTaskInfo> getRunningTasks(int numTasks) {
-        if (mRecentTasks != null
-                && mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_PC)) {
+        if (mRecentTasks != null && shouldEnableRunningTasksForDesktopMode()) {
             try {
                 return new ArrayList<>(Arrays.asList(mRecentTasks.getRunningTasks(numTasks)));
             } catch (RemoteException e) {
@@ -1377,6 +1382,12 @@
         return new ArrayList<>();
     }
 
+    private boolean shouldEnableRunningTasksForDesktopMode() {
+        // TODO(b/335401172): unify DesktopMode checks in Launcher
+        return (enableDesktopWindowingMode() && enableDesktopWindowingTaskbarRunningApps())
+                || mContext.getPackageManager().hasSystemFeature(FEATURE_PC);
+    }
+
     private boolean handleMessageAsync(Message msg) {
         switch (msg.what) {
             case MSG_SET_SHELF_HEIGHT:
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index dec8a12..9d899fc 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -59,6 +59,7 @@
     public static final boolean SHELL_TRANSITIONS_ROTATION = ENABLE_SHELL_TRANSITIONS
             && SystemProperties.getBoolean("persist.wm.debug.shell_transit_rotate", false);
 
+    private final Context mCtx;
     private RecentsAnimationController mController;
     private RecentsAnimationCallbacks mCallbacks;
     private RecentsAnimationTargets mTargets;
@@ -66,7 +67,6 @@
     private GestureState mLastGestureState;
     private RemoteAnimationTarget[] mLastAppearedTaskTargets;
     private Runnable mLiveTileCleanUpHandler;
-    private Context mCtx;
 
     private boolean mRecentsAnimationStartPending = false;
     private boolean mShouldIgnoreMotionEvents = false;
@@ -329,7 +329,7 @@
                 options.setTransientLaunch();
             }
             options.setSourceInfo(ActivityOptions.SourceInfo.TYPE_RECENTS_ANIMATION, eventTime);
-            mRecentsAnimationStartPending = SystemUiProxy.INSTANCE.getNoCreate()
+            mRecentsAnimationStartPending = SystemUiProxy.INSTANCE.get(mCtx)
                     .startRecentsActivity(intent, options, mCallbacks);
             if (enableHandleDelayedGestureCallbacks()) {
                 ActiveGestureLog.INSTANCE.addLog(new ActiveGestureLog.CompoundString(
diff --git a/quickstep/src/com/android/quickstep/TaskUtils.java b/quickstep/src/com/android/quickstep/TaskUtils.java
index 80a449b..63e536a 100644
--- a/quickstep/src/com/android/quickstep/TaskUtils.java
+++ b/quickstep/src/com/android/quickstep/TaskUtils.java
@@ -70,7 +70,7 @@
             return "";
         }
         UserHandle user = UserHandle.of(userId);
-        ApplicationInfo applicationInfo = new PackageManagerHelper(context)
+        ApplicationInfo applicationInfo = PackageManagerHelper.INSTANCE.get(context)
                 .getApplicationInfo(packageName, user, 0);
         if (applicationInfo == null) {
             Log.e(TAG, "Failed to get title for userId=" + userId + ", packageName=" + packageName);
diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index a2a6dde..3a6b804 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -34,6 +34,7 @@
 import androidx.annotation.UiThread;
 
 import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.launcher3.util.SplitConfigurationOptions.SplitStageInfo;
 import com.android.launcher3.util.SplitConfigurationOptions.StagePosition;
@@ -57,7 +58,8 @@
  * This class tracked the top-most task and  some 'approximate' task history to allow faster
  * system state estimation during touch interaction
  */
-public class TopTaskTracker extends ISplitScreenListener.Stub implements TaskStackChangeListener {
+public class TopTaskTracker extends ISplitScreenListener.Stub
+        implements TaskStackChangeListener, SafeCloseable {
 
     public static MainThreadInitializedObject<TopTaskTracker> INSTANCE =
             new MainThreadInitializedObject<>(TopTaskTracker::new);
@@ -67,12 +69,13 @@
     // Ordered list with first item being the most recent task.
     private final LinkedList<RunningTaskInfo> mOrderedTaskList = new LinkedList<>();
 
-
+    private final Context mContext;
     private final SplitStageInfo mMainStagePosition = new SplitStageInfo();
     private final SplitStageInfo mSideStagePosition = new SplitStageInfo();
     private int mPinnedTaskId = INVALID_TASK_ID;
 
     private TopTaskTracker(Context context) {
+        mContext = context;
         mMainStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_MAIN;
         mSideStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_SIDE;
 
@@ -81,6 +84,12 @@
     }
 
     @Override
+    public void close() {
+        TaskStackChangeListeners.getInstance().unregisterTaskStackListener(this);
+        SystemUiProxy.INSTANCE.get(mContext).unregisterSplitScreenListener(this);
+    }
+
+    @Override
     public void onTaskRemoved(int taskId) {
         mOrderedTaskList.removeIf(rto -> rto.taskId == taskId);
     }
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index a09e027..ed633df 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -347,7 +347,6 @@
                         event.getId() + "";
                 Log.d(TAG, name);
             }
-            LauncherAppState appState = LauncherAppState.getInstanceNoCreate();
 
             if (mSlice == null && mSliceItem != null) {
                 mSlice = LauncherAtom.Slice.newBuilder().setUri(
@@ -369,15 +368,10 @@
                 return;
             }
 
-            if (mItemInfo.container < 0 || appState == null) {
-                // 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())));
-            } else {
+            if (mItemInfo.container < 0 || !LauncherAppState.INSTANCE.executeIfCreated(app -> {
                 // Item is inside a collection, fetch collection info in a BG thread
                 // and then write to StatsLog.
-                appState.getModel().enqueueModelUpdateTask(
+                app.getModel().enqueueModelUpdateTask(
                         new BaseModelUpdateTask() {
                             @Override
                             public void execute(@NonNull final LauncherAppState app,
@@ -388,6 +382,11 @@
                                 write(event, applyOverwrites(mItemInfo.buildProto(collectionInfo)));
                             }
                         });
+            })) {
+                // 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())));
             }
         }
 
diff --git a/quickstep/src/com/android/quickstep/util/AppPairsController.java b/quickstep/src/com/android/quickstep/util/AppPairsController.java
index 2b4d280..a82031a 100644
--- a/quickstep/src/com/android/quickstep/util/AppPairsController.java
+++ b/quickstep/src/com/android/quickstep/util/AppPairsController.java
@@ -316,7 +316,7 @@
                 itemInfos.stream().map(ItemInfo::getComponentKey).toList();
 
         // Use TopTaskTracker to find the currently running app (or apps)
-        TopTaskTracker topTaskTracker = getTopTaskTracker(context);
+        TopTaskTracker topTaskTracker = getTopTaskTracker();
 
         // getRunningSplitTasksIds() will return a pair of ids if we are currently running a
         // split pair, or an empty array with zero length if we are running a single app.
@@ -489,7 +489,7 @@
      * Gets the TopTaskTracker, which is a cached record of the top running Task.
      */
     @VisibleForTesting
-    public TopTaskTracker getTopTaskTracker(Context context) {
-        return TopTaskTracker.INSTANCE.get(context);
+    public TopTaskTracker getTopTaskTracker() {
+        return TopTaskTracker.INSTANCE.get(mContext);
     }
 }
diff --git a/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java b/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
index b6e6bf7..2396512 100644
--- a/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
@@ -63,7 +63,7 @@
             SplitSelectStateController controller) {
         mLauncher = launcher;
         mController = controller;
-        mIconCache = LauncherAppState.getInstanceNoCreate().getIconCache();
+        mIconCache = LauncherAppState.getInstance(launcher).getIconCache();
         mHalfDividerSize = mLauncher.getResources().getDimensionPixelSize(
                 R.dimen.multi_window_task_divider_size) / 2;
     }
diff --git a/quickstep/src/com/android/quickstep/util/TaskRemovedDuringLaunchListener.java b/quickstep/src/com/android/quickstep/util/TaskRemovedDuringLaunchListener.java
index 89d8cc4..e80d2a6 100644
--- a/quickstep/src/com/android/quickstep/util/TaskRemovedDuringLaunchListener.java
+++ b/quickstep/src/com/android/quickstep/util/TaskRemovedDuringLaunchListener.java
@@ -22,6 +22,8 @@
 import static com.android.launcher3.BaseActivity.EVENT_RESUMED;
 import static com.android.launcher3.BaseActivity.EVENT_STOPPED;
 
+import android.content.Context;
+
 import androidx.annotation.NonNull;
 
 import com.android.quickstep.RecentsModel;
@@ -45,6 +47,12 @@
     private final Runnable mUnregisterCallback = this::unregister;
     private final Runnable mResumeCallback = this::checkTaskLaunchFailed;
 
+    private final Context mContext;
+
+    public TaskRemovedDuringLaunchListener(Context context) {
+        mContext = context;
+    }
+
     /**
      * Registers a failure listener callback if it detects a scenario in which an app launch
      * failed before the transition finished.
@@ -88,7 +96,7 @@
         if (mLaunchedTaskId != INVALID_TASK_ID) {
             final int launchedTaskId = mLaunchedTaskId;
             final Runnable taskLaunchFailedCallback = mTaskLaunchFailedCallback;
-            RecentsModel.INSTANCE.getNoCreate().isTaskRemoved(mLaunchedTaskId, (taskRemoved) -> {
+            RecentsModel.INSTANCE.get(mContext).isTaskRemoved(mLaunchedTaskId, (taskRemoved) -> {
                 if (taskRemoved) {
                     ActiveGestureLog.INSTANCE.addLog(
                             new ActiveGestureLog.CompoundString("Launch failed, task (id=")
diff --git a/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java b/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
index 8fa5375..82ba30b 100644
--- a/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
+++ b/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
@@ -19,7 +19,7 @@
 import static android.provider.Settings.ACTION_APP_USAGE_SETTINGS;
 
 import static com.android.launcher3.Utilities.prefixTextWithIcon;
-import static com.android.launcher3.util.Executors.THREAD_POOL_EXECUTOR;
+import static com.android.launcher3.util.Executors.ORDERED_BG_EXECUTOR;
 
 import android.app.ActivityOptions;
 import android.content.ActivityNotFoundException;
@@ -140,7 +140,7 @@
     public void initialize(Task task) {
         mAppUsageLimitTimeMs = mAppRemainingTimeMs = -1;
         mTask = task;
-        THREAD_POOL_EXECUTOR.execute(() -> {
+        ORDERED_BG_EXECUTOR.execute(() -> {
                     AppUsageLimit usageLimit = null;
                     try {
                         usageLimit = mLauncherApps.getAppUsageLimit(
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index ae6f703..e910cc3 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -2734,18 +2734,20 @@
      * Returns true if we should add a stub taskView for the running task id
      */
     protected boolean shouldAddStubTaskView(Task[] runningTasks) {
-        if (runningTasks.length > 1) {
-            TaskView primaryTaskView = getTaskViewByTaskId(runningTasks[0].key.id);
-            TaskView secondaryTaskView = getTaskViewByTaskId(runningTasks[1].key.id);
-            int leftTopTaskViewId =
-                    (primaryTaskView == null) ? -1 : primaryTaskView.getTaskViewId();
-            int rightBottomTaskViewId =
-                    (secondaryTaskView == null) ? -1 : secondaryTaskView.getTaskViewId();
-            // Add a new stub view if both taskIds don't match any taskViews
-            return leftTopTaskViewId != rightBottomTaskViewId || leftTopTaskViewId == -1;
+        TaskView taskView = getTaskViewByTaskId(runningTasks[0].key.id);
+        if (taskView == null) {
+            // No TaskView found, add a stub task.
+            return true;
         }
-        Task runningTaskInfo = runningTasks[0];
-        return runningTaskInfo != null && getTaskViewByTaskId(runningTaskInfo.key.id) == null;
+
+        if (runningTasks.length > 1) {
+            // Ensure all taskIds matches the TaskView, otherwise add a stub task.
+            return Arrays.stream(runningTasks).anyMatch(
+                    runningTask -> !taskView.containsTaskId(runningTask.key.id));
+        } else {
+            // Ensure the TaskView only contains a single taskId, otherwise add a stub task.
+            return taskView.containsMultipleTasks();
+        }
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 9a850b5..0c7c5a5 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -91,6 +91,7 @@
 import com.android.launcher3.util.CancellableTask;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.RunnableList;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.launcher3.util.SplitConfigurationOptions.SplitPositionOption;
 import com.android.launcher3.util.TraceHelper;
@@ -915,8 +916,8 @@
             TestLogging.recordEvent(
                     TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
 
-            final TaskRemovedDuringLaunchListener
-                    failureListener = new TaskRemovedDuringLaunchListener();
+            TaskRemovedDuringLaunchListener failureListener = new TaskRemovedDuringLaunchListener(
+                    getContext().getApplicationContext());
             if (isQuickswitch) {
                 // We only listen for failures to launch in quickswitch because the during this
                 // gesture launcher is in the background state, vs other launches which are in
@@ -943,12 +944,8 @@
             // Indicate success once the system has indicated that the transition has started
             ActivityOptions opts = ActivityOptions.makeCustomTaskAnimation(getContext(), 0, 0,
                     MAIN_EXECUTOR.getHandler(),
-                    elapsedRealTime -> {
-                        callback.accept(true);
-                    },
-                    elapsedRealTime -> {
-                        failureListener.onTransitionFinished();
-                    });
+                    elapsedRealTime -> callback.accept(true),
+                    elapsedRealTime -> failureListener.onTransitionFinished());
             opts.setLaunchDisplayId(
                     getDisplay() == null ? DEFAULT_DISPLAY : getDisplay().getDisplayId());
             if (isQuickswitch) {
@@ -1845,7 +1842,7 @@
     /**
      * We update and subsequently draw these in {@link #setFullscreenProgress(float)}.
      */
-    public static class FullscreenDrawParams {
+    public static class FullscreenDrawParams implements SafeCloseable {
 
         private float mCornerRadius;
         private float mWindowCornerRadius;
@@ -1880,6 +1877,9 @@
                     Utilities.mapRange(fullscreenProgress, mCornerRadius, mWindowCornerRadius)
                             / parentScale / taskViewScale;
         }
+
+        @Override
+        public void close() { }
     }
 
     public class TaskIdAttributeContainer {
diff --git a/quickstep/tests/OWNERS b/quickstep/tests/OWNERS
index c271803..02e8ebc 100644
--- a/quickstep/tests/OWNERS
+++ b/quickstep/tests/OWNERS
@@ -2,4 +2,3 @@
 sunnygoyal@google.com
 winsonc@google.com
 hyunyoungs@google.com
-mateuszc@google.com
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
index 9fa4b79..72cfd92 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
@@ -44,7 +44,6 @@
 import com.android.launcher3.util.window.CachedDisplayInfo;
 import com.android.launcher3.util.window.WindowManagerProxy;
 import com.android.quickstep.FallbackActivityInterface;
-import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.util.SurfaceTransaction.MockProperties;
 
 import org.hamcrest.Description;
@@ -160,7 +159,6 @@
         void verifyNoTransforms() {
             LauncherModelHelper helper = new LauncherModelHelper();
             try {
-                helper.sandboxContext.allow(SystemUiProxy.INSTANCE);
                 int rotation = mDisplaySize.x > mDisplaySize.y
                         ? Surface.ROTATION_90 : Surface.ROTATION_0;
                 CachedDisplayInfo cdi = new CachedDisplayInfo(mDisplaySize, rotation);
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/DesktopTaskbarRunningAppsControllerTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/DesktopTaskbarRunningAppsControllerTest.kt
new file mode 100644
index 0000000..93eefe2
--- /dev/null
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/DesktopTaskbarRunningAppsControllerTest.kt
@@ -0,0 +1,205 @@
+/*
+ * 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
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.content.ComponentName
+import android.content.Intent
+import android.os.Process
+import android.os.UserHandle
+import android.testing.AndroidTestingRunner
+import android.util.SparseArray
+import com.android.launcher3.model.data.AppInfo
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.statehandlers.DesktopVisibilityController
+import com.android.quickstep.RecentsModel
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.junit.MockitoJUnit
+import org.mockito.kotlin.whenever
+
+@RunWith(AndroidTestingRunner::class)
+class DesktopTaskbarRunningAppsControllerTest : TaskbarBaseTestCase() {
+
+    @get:Rule val mockitoRule = MockitoJUnit.rule()
+
+    @Mock private lateinit var mockRecentsModel: RecentsModel
+    @Mock private lateinit var mockDesktopVisibilityController: DesktopVisibilityController
+
+    private var nextTaskId: Int = 500
+
+    private lateinit var taskbarRunningAppsController: DesktopTaskbarRunningAppsController
+    private lateinit var userHandle: UserHandle
+
+    @Before
+    fun setUp() {
+        super.setup()
+        userHandle = Process.myUserHandle()
+        taskbarRunningAppsController =
+            DesktopTaskbarRunningAppsController(mockRecentsModel, mockDesktopVisibilityController)
+        taskbarRunningAppsController.init(taskbarControllers)
+        taskbarRunningAppsController.setApps(
+            ALL_APP_PACKAGES.map { createTestAppInfo(packageName = it) }.toTypedArray()
+        )
+    }
+
+    @Test
+    fun updateHotseatItemInfos_null_returnsNull() {
+        assertThat(taskbarRunningAppsController.updateHotseatItemInfos(/* hotseatItems= */ null))
+            .isNull()
+    }
+
+    @Test
+    fun updateHotseatItemInfos_notInDesktopMode_returnsExistingHotseatItems() {
+        setInDesktopMode(false)
+        val hotseatItems =
+            createHotseatItemsFromPackageNames(listOf(HOTSEAT_PACKAGE_1, HOTSEAT_PACKAGE_2))
+                .toTypedArray()
+
+        assertThat(taskbarRunningAppsController.updateHotseatItemInfos(hotseatItems))
+            .isEqualTo(hotseatItems)
+    }
+
+    @Test
+    fun updateHotseatItemInfos_notInDesktopMode_runningApps_returnsExistingHotseatItems() {
+        setInDesktopMode(false)
+        val hotseatPackages = listOf(HOTSEAT_PACKAGE_1, HOTSEAT_PACKAGE_2)
+        val hotseatItems = createHotseatItemsFromPackageNames(hotseatPackages)
+        val runningTasks =
+            createDesktopTasksFromPackageNames(listOf(RUNNING_APP_PACKAGE_1, RUNNING_APP_PACKAGE_2))
+        whenever(mockRecentsModel.runningTasks).thenReturn(runningTasks)
+        taskbarRunningAppsController.updateRunningApps(createSparseArray(hotseatItems))
+
+        val newHotseatItems =
+            taskbarRunningAppsController.updateHotseatItemInfos(hotseatItems.toTypedArray())
+
+        assertThat(newHotseatItems?.map { it.targetPackage }).isEqualTo(hotseatPackages)
+    }
+
+    @Test
+    fun updateHotseatItemInfos_noRunningApps_returnsExistingHotseatItems() {
+        setInDesktopMode(true)
+        val hotseatItems: Array<ItemInfo> =
+            createHotseatItemsFromPackageNames(listOf(HOTSEAT_PACKAGE_1, HOTSEAT_PACKAGE_2))
+                .toTypedArray()
+
+        assertThat(taskbarRunningAppsController.updateHotseatItemInfos(hotseatItems))
+            .isEqualTo(hotseatItems)
+    }
+
+    @Test
+    fun updateHotseatItemInfos_returnsExistingHotseatItemsAndRunningApps() {
+        setInDesktopMode(true)
+        val hotseatItems =
+            createHotseatItemsFromPackageNames(listOf(HOTSEAT_PACKAGE_1, HOTSEAT_PACKAGE_2))
+        val runningTasks =
+            createDesktopTasksFromPackageNames(listOf(RUNNING_APP_PACKAGE_1, RUNNING_APP_PACKAGE_2))
+        whenever(mockRecentsModel.runningTasks).thenReturn(runningTasks)
+        taskbarRunningAppsController.updateRunningApps(createSparseArray(hotseatItems))
+
+        val newHotseatItems =
+            taskbarRunningAppsController.updateHotseatItemInfos(hotseatItems.toTypedArray())
+
+        val expectedPackages =
+            listOf(
+                HOTSEAT_PACKAGE_1,
+                HOTSEAT_PACKAGE_2,
+                RUNNING_APP_PACKAGE_1,
+                RUNNING_APP_PACKAGE_2,
+            )
+        assertThat(newHotseatItems?.map { it.targetPackage }).isEqualTo(expectedPackages)
+    }
+
+    @Test
+    fun updateHotseatItemInfos_runningAppIsHotseatItem_returnsDistinctItems() {
+        setInDesktopMode(true)
+        val hotseatItems =
+            createHotseatItemsFromPackageNames(listOf(HOTSEAT_PACKAGE_1, HOTSEAT_PACKAGE_2))
+        val runningTasks =
+            createDesktopTasksFromPackageNames(
+                listOf(HOTSEAT_PACKAGE_1, RUNNING_APP_PACKAGE_1, RUNNING_APP_PACKAGE_2)
+            )
+        whenever(mockRecentsModel.runningTasks).thenReturn(runningTasks)
+        taskbarRunningAppsController.updateRunningApps(createSparseArray(hotseatItems))
+
+        val newHotseatItems =
+            taskbarRunningAppsController.updateHotseatItemInfos(hotseatItems.toTypedArray())
+
+        val expectedPackages =
+            listOf(
+                HOTSEAT_PACKAGE_1,
+                HOTSEAT_PACKAGE_2,
+                RUNNING_APP_PACKAGE_1,
+                RUNNING_APP_PACKAGE_2,
+            )
+        assertThat(newHotseatItems?.map { it.targetPackage }).isEqualTo(expectedPackages)
+    }
+
+    private fun createHotseatItemsFromPackageNames(packageNames: List<String>): List<ItemInfo> {
+        return packageNames.map { createTestAppInfo(packageName = it) }
+    }
+
+    private fun createDesktopTasksFromPackageNames(
+        packageNames: List<String>
+    ): ArrayList<RunningTaskInfo> {
+        return ArrayList(packageNames.map { createDesktopTaskInfo(packageName = it) })
+    }
+
+    private fun createDesktopTaskInfo(packageName: String): RunningTaskInfo {
+        return RunningTaskInfo().apply {
+            taskId = nextTaskId++
+            configuration.windowConfiguration.windowingMode = WINDOWING_MODE_FREEFORM
+            realActivity = ComponentName(packageName, "TestActivity")
+        }
+    }
+
+    private fun createTestAppInfo(
+        packageName: String = "testPackageName",
+        className: String = "testClassName"
+    ) = AppInfo(ComponentName(packageName, className), className /* title */, userHandle, Intent())
+
+    private fun setInDesktopMode(inDesktopMode: Boolean) {
+        whenever(mockDesktopVisibilityController.areDesktopTasksVisible()).thenReturn(inDesktopMode)
+    }
+
+    private fun createSparseArray(itemInfos: List<ItemInfo>): SparseArray<ItemInfo> {
+        val sparseArray = SparseArray<ItemInfo>()
+        itemInfos.forEachIndexed { index, itemInfo -> sparseArray[index] = itemInfo }
+        return sparseArray
+    }
+
+    private companion object {
+        const val HOTSEAT_PACKAGE_1 = "hotseat1"
+        const val HOTSEAT_PACKAGE_2 = "hotseat2"
+        const val RUNNING_APP_PACKAGE_1 = "running1"
+        const val RUNNING_APP_PACKAGE_2 = "running2"
+        const val RUNNING_APP_PACKAGE_3 = "running3"
+        val ALL_APP_PACKAGES =
+            listOf(
+                HOTSEAT_PACKAGE_1,
+                HOTSEAT_PACKAGE_2,
+                RUNNING_APP_PACKAGE_1,
+                RUNNING_APP_PACKAGE_2,
+                RUNNING_APP_PACKAGE_3,
+            )
+    }
+}
diff --git a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
index 857a4e3..edc0a6f 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
@@ -21,7 +21,6 @@
 
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
-import static org.junit.Assume.assumeTrue;
 
 import android.util.Log;
 
@@ -29,6 +28,8 @@
 import androidx.test.filters.LargeTest;
 
 import com.android.launcher3.tapl.LauncherInstrumentation;
+import com.android.launcher3.tapl.PrivateSpaceContainer;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ScreenRecordRule;
 
 import org.junit.After;
@@ -43,7 +44,9 @@
 public class TaplPrivateSpaceTest extends AbstractQuickStepTest {
 
     private int mProfileUserId;
-    private boolean mPrivateProfileSetupSuccessful;
+
+    private static final String PRIVATE_PROFILE_NAME = "LauncherPrivateProfile";
+    private static final String INSTALLED_APP_NAME = "Aardwolf";
     private static final String TAG = "TaplPrivateSpaceTest";
 
     @Override
@@ -52,8 +55,6 @@
         initialize(this);
 
         createAndStartPrivateProfileUser();
-        assumeTrue("Private Profile Setup not successful, aborting",
-                mPrivateProfileSetupSuccessful);
 
         mDevice.pressHome();
         waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
@@ -72,7 +73,7 @@
 
     private void createAndStartPrivateProfileUser() {
         String createUserOutput = executeShellCommand("pm create-user --profileOf 0 --user-type "
-                + "android.os.usertype.profile.PRIVATE LauncherPrivateProfile");
+                + "android.os.usertype.profile.PRIVATE " + PRIVATE_PROFILE_NAME);
         updatePrivateProfileSetupSuccessful("pm create-user", createUserOutput);
         String[] tokens = createUserOutput.split("\\s+");
         mProfileUserId = Integer.parseInt(tokens[tokens.length - 1]);
@@ -86,24 +87,42 @@
 
     @After
     public void removePrivateProfile() {
-        String output = executeShellCommand("pm remove-user " + mProfileUserId);
-        updateProfileRemovalSuccessful("pm remove-user", output);
-        waitForPrivateSpaceRemoval();
+        String userListOutput = executeShellCommand("pm list users");
+        if (isPrivateProfilePresent("pm list users", userListOutput)) {
+            String output = executeShellCommand("pm remove-user " + mProfileUserId);
+            updateProfileRemovalSuccessful("pm remove-user", output);
+            waitForPrivateSpaceRemoval();
+        }
     }
 
     @Test
     @ScreenRecordRule.ScreenRecord // b/334946529
     public void testPrivateSpaceContainerIsPresent() {
-        assumeTrue(mPrivateProfileSetupSuccessful);
         // Scroll to the bottom of All Apps
         executeOnLauncher(launcher -> launcher.getAppsView().resetAndScrollToPrivateSpaceHeader());
-        waitForResumed("Launcher internal state is still Background");
 
         // Verify Unlocked View elements are present.
         assertNotNull("Private Space Unlocked View not found, or is not correct",
                 mLauncher.getAllApps().getPrivateSpaceUnlockedView());
     }
 
+    @Test
+    @ScreenRecordRule.ScreenRecord // b/334946529
+    public void testUserInstalledAppIsShownAboveDivider() throws IOException {
+        // Ensure that the App is not installed in main user otherwise, it may not be found in
+        // PS container.
+        TestUtil.uninstallDummyApp();
+        // Install the app in Private Profile
+        TestUtil.installDummyAppForUser(mProfileUserId);
+        waitForLauncherUIUpdate();
+        // Scroll to the bottom of All Apps
+        executeOnLauncher(launcher -> launcher.getAppsView().resetAndScrollToPrivateSpaceHeader());
+
+        // Verify the Installed App is displayed in correct position.
+        PrivateSpaceContainer psContainer = mLauncher.getAllApps().getPrivateSpaceUnlockedView();
+        psContainer.verifyInstalledAppIsPresent(INSTALLED_APP_NAME);
+    }
+
     private void waitForPrivateSpaceSetup() {
         waitForLauncherCondition("Private Profile not setup",
                 launcher -> launcher.getAppsView().hasPrivateProfile(),
@@ -129,7 +148,7 @@
     private void updatePrivateProfileSetupSuccessful(String cli, String output) {
         Log.d(TAG, "updatePrivateProfileSetupSuccessful, cli=" + cli + " " + "output="
                 + output);
-        mPrivateProfileSetupSuccessful = output.startsWith("Success");
+        assertTrue(output, output.startsWith("Success"));
     }
 
     private void updateProfileRemovalSuccessful(String cli, String output) {
@@ -137,6 +156,11 @@
         assertTrue(output, output.startsWith("Success"));
     }
 
+    private boolean isPrivateProfilePresent(String cli, String output) {
+        Log.d(TAG, "updatePrivateProfilePresent, cli=" + cli + " " + "output=" + output);
+        return output.contains(PRIVATE_PROFILE_NAME);
+    }
+
     private String executeShellCommand(String command) {
         try {
             return mDevice.executeShellCommand(command);
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsKeyboardQuickSwitch.java b/quickstep/tests/src/com/android/quickstep/TaplTestsKeyboardQuickSwitch.java
index 43ebb17..7ab72f2 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsKeyboardQuickSwitch.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsKeyboardQuickSwitch.java
@@ -65,7 +65,7 @@
     public void setUp() throws Exception {
         Assume.assumeTrue(mLauncher.isTablet());
         super.setUp();
-        startAppFast(CALCULATOR_APP_PACKAGE);
+        startAppFastInFullscreen(CALCULATOR_APP_PACKAGE);
         startTestActivity(2);
     }
 
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java b/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
index f0683f9..ec245ee 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
@@ -15,6 +15,8 @@
  */
 package com.android.quickstep;
 
+import static androidx.test.InstrumentationRegistry.getTargetContext;
+
 import static com.android.launcher3.util.TestConstants.AppNames.TEST_APP_NAME;
 import static com.android.quickstep.TaplTestsTaskbar.TaskbarMode.PERSISTENT;
 import static com.android.quickstep.TaplTestsTaskbar.TaskbarMode.TRANSIENT;
@@ -53,7 +55,7 @@
 
     @Override
     public void setUp() throws Exception {
-        mTaskbarWasInTransientMode = isTaskbarInTransientMode(mTargetContext);
+        mTaskbarWasInTransientMode = isTaskbarInTransientMode(getTargetContext());
         setTaskbarMode(mLauncher, isTaskbarTestModeTransient());
         super.setUp();
     }
diff --git a/quickstep/tests/src/com/android/quickstep/util/AppPairsControllerTest.kt b/quickstep/tests/src/com/android/quickstep/util/AppPairsControllerTest.kt
index adaf7ff..ece67af 100644
--- a/quickstep/tests/src/com/android/quickstep/util/AppPairsControllerTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/util/AppPairsControllerTest.kt
@@ -43,6 +43,7 @@
 import org.mockito.kotlin.any
 import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.doNothing
+import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.never
 import org.mockito.kotlin.spy
@@ -100,8 +101,7 @@
         // Stub methods on appPairsController so that they return mocks
         spyAppPairsController = spy(appPairsController)
         whenever(mockAppPairIcon.context).thenReturn(mockTaskbarActivityContext)
-        whenever(spyAppPairsController.getTopTaskTracker(mockTaskbarActivityContext))
-            .thenReturn(mockTopTaskTracker)
+        doReturn(mockTopTaskTracker).whenever(spyAppPairsController).topTaskTracker
         whenever(mockTopTaskTracker.getCachedTopTask(any())).thenReturn(mockCachedTaskInfo)
         whenever(mockTask1.getKey()).thenReturn(mockTaskKey1)
         whenever(mockTask2.getKey()).thenReturn(mockTaskKey2)
diff --git a/res/drawable/ic_private_profile_app_scroller_badge.xml b/res/drawable/ic_private_profile_app_scroller_badge.xml
index b52a277..ede42b9 100644
--- a/res/drawable/ic_private_profile_app_scroller_badge.xml
+++ b/res/drawable/ic_private_profile_app_scroller_badge.xml
@@ -21,8 +21,8 @@
     <path
         android:pathData="M16.0007 2.66602L5.33398 6.66602V14.786C5.33398 21.5194 9.88065 27.7993 16.0007 29.3327C22.1207 27.7993 26.6673 21.5194 26.6673 14.786V6.66602L16.0007 2.66602ZM20.0007 19.9993V22.666H17.334V23.9993H14.6673V17.1193C12.7473 16.546 11.334 14.786 11.334 12.666C11.334 10.0927 13.4273 7.99935 16.0007 7.99935C18.574 7.99935 20.6673 10.0927 20.6673 12.666C20.6673 14.7727 19.254 16.546 17.334 17.1193V19.9993H20.0007Z"
         android:fillType="evenOdd"
-        android:fillColor="@android:color/white" />
+        android:fillColor="?android:attr/textColorPrimaryInverse" />
     <path
         android:pathData="M16 14.666C17.1046 14.666 18 13.7706 18 12.666C18 11.5614 17.1046 10.666 16 10.666C14.8954 10.666 14 11.5614 14 12.666C14 13.7706 14.8954 14.666 16 14.666Z"
-        android:fillColor="@android:color/white" />
+        android:fillColor="?android:attr/textColorPrimaryInverse" />
 </vector>
diff --git a/res/layout/private_space_header.xml b/res/layout/private_space_header.xml
index 65f1004..84554e4 100644
--- a/res/layout/private_space_header.xml
+++ b/res/layout/private_space_header.xml
@@ -40,7 +40,6 @@
             android:layout_width="@dimen/ps_header_image_height"
             android:layout_height="@dimen/ps_header_image_height"
             android:background="@drawable/ps_settings_background"
-            android:layout_marginEnd="@dimen/ps_header_settings_icon_margin_end"
             android:src="@drawable/ic_ps_settings"
             android:contentDescription="@string/ps_container_settings" />
         <LinearLayout
@@ -49,7 +48,7 @@
             android:layout_height="@dimen/ps_header_image_height"
             android:background="@drawable/ps_lock_background"
             android:gravity="center_vertical"
-            android:layout_marginEnd="@dimen/ps_header_layout_margin"
+            android:layout_marginEnd="@dimen/ps_lock_button_margin_end"
             android:contentDescription="@string/ps_container_lock_unlock_button">
             <ImageView
                 android:id="@+id/lock_icon"
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 5ff9902..5dd3569 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -496,7 +496,7 @@
     <dimen name="ps_header_image_height">48dp</dimen>
     <dimen name="ps_header_text_height">24dp</dimen>
     <dimen name="ps_header_layout_margin">16dp</dimen>
-    <dimen name="ps_header_settings_icon_margin_end">4dp</dimen>
+    <dimen name="ps_lock_button_margin_end">12dp</dimen>
     <dimen name="ps_header_text_size">16sp</dimen>
     <dimen name="ps_button_height">40dp</dimen>
     <dimen name="ps_button_width">40dp</dimen>
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 98cb84e..f405b93 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -59,6 +59,7 @@
 import com.android.launcher3.util.LockedUserState;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Partner;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.util.window.WindowManagerProxy;
 
@@ -75,7 +76,7 @@
 import java.util.Objects;
 import java.util.stream.Collectors;
 
-public class InvariantDeviceProfile {
+public class InvariantDeviceProfile implements SafeCloseable {
 
     public static final String TAG = "IDP";
     // We do not need any synchronization for this variable as its only written on UI thread.
@@ -229,9 +230,8 @@
         if (!newGridName.equals(gridName)) {
             LauncherPrefs.get(context).put(GRID_NAME, newGridName);
         }
-        LockedUserState.get(context).runOnUserUnlocked(() -> {
-            new DeviceGridState(this).writeToPrefs(context);
-        });
+        LockedUserState.get(context).runOnUserUnlocked(() ->
+            new DeviceGridState(this).writeToPrefs(context));
 
         DisplayController.INSTANCE.get(context).setPriorityListener(
                 (displayContext, info, flags) -> {
@@ -295,6 +295,11 @@
         initGrid(context, myInfo, result, deviceType);
     }
 
+    @Override
+    public void close() {
+        DisplayController.INSTANCE.executeIfCreated(dc -> dc.setPriorityListener(null));
+    }
+
     /**
      * Reinitialize the current grid after a restore, where some grids might now be disabled.
      */
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 50a597d..d2633e0 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -78,14 +78,10 @@
 
     private final RunnableList mOnTerminateCallback = new RunnableList();
 
-    public static LauncherAppState getInstance(final Context context) {
+    public static LauncherAppState getInstance(Context context) {
         return INSTANCE.get(context);
     }
 
-    public static LauncherAppState getInstanceNoCreate() {
-        return INSTANCE.getNoCreate();
-    }
-
     public Context getContext() {
         return mContext;
     }
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 7fdfd72..9b0e0ec 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -447,8 +447,8 @@
                 IconCache iconCache = app.getIconCache();
                 final IntSet removedIds = new IntSet();
                 HashSet<WorkspaceItemInfo> archivedWorkspaceItemsToCacheRefresh = new HashSet<>();
-                boolean isAppArchived = new PackageManagerHelper(
-                        mApp.getContext()).isAppArchivedForUser(packageName, user);
+                boolean isAppArchived = PackageManagerHelper.INSTANCE.get(mApp.getContext())
+                        .isAppArchivedForUser(packageName, user);
                 synchronized (dataModel) {
                     if (isAppArchived) {
                         // Remove package icon cache entry for archived app in case of a session
diff --git a/src/com/android/launcher3/LauncherPrefs.kt b/src/com/android/launcher3/LauncherPrefs.kt
index 80a8cea..b503739 100644
--- a/src/com/android/launcher3/LauncherPrefs.kt
+++ b/src/com/android/launcher3/LauncherPrefs.kt
@@ -30,6 +30,7 @@
 import com.android.launcher3.states.RotationHelper
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.MainThreadInitializedObject
+import com.android.launcher3.util.SafeCloseable
 import com.android.launcher3.util.Themes
 
 /**
@@ -37,7 +38,7 @@
  *
  * TODO(b/262721340): Replace all direct SharedPreference refs with LauncherPrefs / Item methods.
  */
-class LauncherPrefs(private val encryptedContext: Context) {
+class LauncherPrefs(private val encryptedContext: Context) : SafeCloseable {
     private val deviceProtectedStorageContext =
         encryptedContext.createDeviceProtectedStorageContext()
 
@@ -242,6 +243,8 @@
         }
     }
 
+    override fun close() {}
+
     companion object {
         @VisibleForTesting const val BOOT_AWARE_PREFS_KEY = "boot_aware_prefs"
 
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 4e0ba62..6e2d357 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -48,11 +48,11 @@
      */
     @Override
     public void dump(FileDescriptor fd, PrintWriter writer, String[] args) {
-        LauncherAppState appState = LauncherAppState.getInstanceNoCreate();
-        if (appState == null || !appState.getModel().isModelLoaded()) {
-            return;
-        }
-        appState.getModel().dumpState("", fd, writer, args);
+        LauncherAppState.INSTANCE.executeIfCreated(appState -> {
+            if (appState.getModel().isModelLoaded()) {
+                appState.getModel().dumpState("", fd, writer, args);
+            }
+        });
     }
 
     @Override
diff --git a/src/com/android/launcher3/SecondaryDropTarget.java b/src/com/android/launcher3/SecondaryDropTarget.java
index 1362586..0a4fb73 100644
--- a/src/com/android/launcher3/SecondaryDropTarget.java
+++ b/src/com/android/launcher3/SecondaryDropTarget.java
@@ -362,7 +362,7 @@
 
         public void onLauncherResume() {
             // We use MATCH_UNINSTALLED_PACKAGES as the app can be on SD card as well.
-            if (new PackageManagerHelper(mContext).getApplicationInfo(mPackageName,
+            if (PackageManagerHelper.INSTANCE.get(mContext).getApplicationInfo(mPackageName,
                     mDragObject.dragInfo.user, PackageManager.MATCH_UNINSTALLED_PACKAGES) == null) {
                 mDragObject.dragSource = mOriginal;
                 mOriginal.onDropCompleted(SecondaryDropTarget.this, mDragObject, true);
diff --git a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
index 60df7c5..a810331 100644
--- a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
+++ b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
@@ -22,7 +22,6 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_USER_INSTALLED_APPS_COUNT;
 
 import android.content.Context;
-import android.graphics.drawable.Drawable;
 import android.text.Spannable;
 import android.text.SpannableString;
 import android.text.style.ImageSpan;
@@ -362,7 +361,7 @@
         // Split of private space apps into user-installed and system apps.
         Map<Boolean, List<AppInfo>> split = mPrivateApps.stream()
                 .collect(Collectors.partitioningBy(mPrivateProviderManager
-                                .splitIntoUserInstalledAndSystemApps()));
+                                .splitIntoUserInstalledAndSystemApps(mActivityContext)));
 
         // TODO(b/329688630): switch to the pulled LayoutStaticSnapshot atom
         mActivityContext
diff --git a/src/com/android/launcher3/allapps/PrivateProfileManager.java b/src/com/android/launcher3/allapps/PrivateProfileManager.java
index 38fe138..7e975df 100644
--- a/src/com/android/launcher3/allapps/PrivateProfileManager.java
+++ b/src/com/android/launcher3/allapps/PrivateProfileManager.java
@@ -80,9 +80,7 @@
 import com.android.launcher3.views.RecyclerViewFastScroller;
 
 import java.util.ArrayList;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 import java.util.function.Predicate;
 
 /**
@@ -105,7 +103,6 @@
             }
         }
     };
-    private Set<String> mPreInstalledSystemPackages = new HashSet<>();
     private Intent mAppInstallerIntent = new Intent();
     private PrivateAppsSectionDecorator mPrivateAppsSectionDecorator;
     private boolean mPrivateSpaceSettingsAvailable;
@@ -264,8 +261,6 @@
         ApiWrapper apiWrapper = ApiWrapper.INSTANCE.get(appContext);
         UserHandle profileUser = getProfileUser();
         if (profileUser != null) {
-            mPreInstalledSystemPackages = new HashSet<>(
-                    apiWrapper.getPreInstalledSystemPackages(profileUser));
             mAppInstallerIntent = apiWrapper
                     .getAppMarketActivityIntent(BuildConfig.APPLICATION_ID, profileUser);
         }
@@ -349,10 +344,12 @@
      * Splits private apps into user installed and system apps.
      * When the list of system apps is empty, all apps are treated as system.
      */
-    public Predicate<AppInfo> splitIntoUserInstalledAndSystemApps() {
-        return appInfo -> !mPreInstalledSystemPackages.isEmpty()
+    public Predicate<AppInfo> splitIntoUserInstalledAndSystemApps(Context context) {
+        List<String> preInstallApps = UserCache.getInstance(context)
+                .getPreInstallApps(getProfileUser());
+        return appInfo -> !preInstallApps.isEmpty()
                 && (appInfo.componentName == null
-                || !(mPreInstalledSystemPackages.contains(appInfo.componentName.getPackageName())));
+                || !(preInstallApps.contains(appInfo.componentName.getPackageName())));
     }
 
     /** Add Private Space Header view elements based upon {@link UserProfileState} */
diff --git a/src/com/android/launcher3/dragndrop/AddItemActivity.java b/src/com/android/launcher3/dragndrop/AddItemActivity.java
index 05fdcef..ec0a222 100644
--- a/src/com/android/launcher3/dragndrop/AddItemActivity.java
+++ b/src/com/android/launcher3/dragndrop/AddItemActivity.java
@@ -162,7 +162,7 @@
             finish();
             return;
         }
-        ApplicationInfo info = new PackageManagerHelper(this)
+        ApplicationInfo info = PackageManagerHelper.INSTANCE.get(this)
                 .getApplicationInfo(targetApp.packageName, targetApp.user, 0);
         if (info == null) {
             finish();
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 6bed9dc..50d2b43 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -66,7 +66,6 @@
 import com.android.launcher3.InsettableFrameLayout;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
@@ -88,14 +87,11 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.pm.InstallSessionHelper;
-import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
-import com.android.launcher3.util.PluginManagerWrapper;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.util.window.WindowManagerProxy;
 import com.android.launcher3.views.ActivityContext;
@@ -104,7 +100,6 @@
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 import com.android.launcher3.widget.LocalColorExtractor;
-import com.android.launcher3.widget.custom.CustomWidgetManager;
 import com.android.launcher3.widget.util.WidgetSizes;
 
 import java.util.ArrayList;
@@ -136,13 +131,10 @@
                 new ConcurrentLinkedQueue<>();
 
         public PreviewContext(Context base, InvariantDeviceProfile idp) {
-            super(base, UserCache.INSTANCE, InstallSessionHelper.INSTANCE, LauncherPrefs.INSTANCE,
-                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
-                    CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE,
-                    WindowManagerProxy.INSTANCE, DisplayController.INSTANCE);
+            super(base);
             mIdp = idp;
-            mObjectMap.put(InvariantDeviceProfile.INSTANCE, idp);
-            mObjectMap.put(LauncherAppState.INSTANCE,
+            putObject(InvariantDeviceProfile.INSTANCE, idp);
+            putObject(LauncherAppState.INSTANCE,
                     new LauncherAppState(this, null /* iconCacheFileName */));
         }
 
diff --git a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
index ce563b7..3fa6da4 100644
--- a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
+++ b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
@@ -15,6 +15,7 @@
  */
 package com.android.launcher3.model;
 
+import android.content.Context;
 import android.content.Intent;
 import android.content.pm.LauncherActivityInfo;
 import android.content.pm.LauncherApps;
@@ -85,6 +86,7 @@
 
         final ArrayList<ItemInfo> addedItemsFinal = new ArrayList<>();
         final IntArray addedWorkspaceScreensFinal = new IntArray();
+        final Context context = app.getContext();
 
         synchronized (dataModel) {
             IntArray workspaceScreens = dataModel.collectWorkspaceScreens();
@@ -99,7 +101,7 @@
                     }
 
                     // b/139663018 Short-circuit this logic if the icon is a system app
-                    if (PackageManagerHelper.isSystemApp(app.getContext(),
+                    if (PackageManagerHelper.isSystemApp(context,
                             Objects.requireNonNull(item.getIntent()))) {
                         continue;
                     }
@@ -112,7 +114,7 @@
 
                 if (item.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPLICATION) {
                     if (item instanceof WorkspaceItemFactory) {
-                        item = ((WorkspaceItemFactory) item).makeWorkspaceItem(app.getContext());
+                        item = ((WorkspaceItemFactory) item).makeWorkspaceItem(context);
                     }
                 }
                 if (item != null) {
@@ -121,8 +123,8 @@
             }
 
             InstallSessionHelper packageInstaller =
-                    InstallSessionHelper.INSTANCE.get(app.getContext());
-            LauncherApps launcherApps = app.getContext().getSystemService(LauncherApps.class);
+                    InstallSessionHelper.INSTANCE.get(context);
+            LauncherApps launcherApps = context.getSystemService(LauncherApps.class);
 
             for (ItemInfo item : filteredItems) {
                 // Find appropriate space for the item.
@@ -135,7 +137,7 @@
                         || item instanceof LauncherAppWidgetInfo) {
                     itemInfo = item;
                 } else if (item instanceof WorkspaceItemFactory) {
-                    itemInfo = ((WorkspaceItemFactory) item).makeWorkspaceItem(app.getContext());
+                    itemInfo = ((WorkspaceItemFactory) item).makeWorkspaceItem(context);
                 } else {
                     throw new RuntimeException("Unexpected info type");
                 }
diff --git a/src/com/android/launcher3/model/AllAppsList.java b/src/com/android/launcher3/model/AllAppsList.java
index a1a05f4..39c1243 100644
--- a/src/com/android/launcher3/model/AllAppsList.java
+++ b/src/com/android/launcher3/model/AllAppsList.java
@@ -169,7 +169,7 @@
     public AppInfo addPromiseApp(
             Context context, PackageInstallInfo installInfo, boolean loadIcon) {
         // only if not yet installed
-        if (new PackageManagerHelper(context)
+        if (PackageManagerHelper.INSTANCE.get(context)
                 .isAppInstalled(installInfo.packageName, installInfo.user)) {
             return null;
         }
diff --git a/src/com/android/launcher3/model/ItemInstallQueue.java b/src/com/android/launcher3/model/ItemInstallQueue.java
index 90aba2a..551c2d8 100644
--- a/src/com/android/launcher3/model/ItemInstallQueue.java
+++ b/src/com/android/launcher3/model/ItemInstallQueue.java
@@ -45,7 +45,6 @@
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings.Favorites;
-import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
@@ -55,6 +54,7 @@
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.PersistedItemArray;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 
 import java.util.HashSet;
@@ -65,7 +65,7 @@
 /**
  * Class to maintain a queue of pending items to be added to the workspace.
  */
-public class ItemInstallQueue {
+public class ItemInstallQueue implements SafeCloseable {
 
     private static final String LOG = "ItemInstallQueue";
 
@@ -99,6 +99,9 @@
         mContext = context;
     }
 
+    @Override
+    public void close() {}
+
     @WorkerThread
     private void ensureQueueLoaded() {
         Preconditions.assertWorkerThread();
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index a742c75..77bc32e 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -138,6 +138,7 @@
     private final LauncherApps mLauncherApps;
     private final UserManager mUserManager;
     private final UserCache mUserCache;
+    private final PackageManagerHelper mPmHelper;
 
     private final InstallSessionHelper mSessionHelper;
     private final IconCache mIconCache;
@@ -170,6 +171,7 @@
         mLauncherApps = mApp.getContext().getSystemService(LauncherApps.class);
         mUserManager = mApp.getContext().getSystemService(UserManager.class);
         mUserCache = UserCache.INSTANCE.get(mApp.getContext());
+        mPmHelper = PackageManagerHelper.INSTANCE.get(mApp.getContext());
         mSessionHelper = InstallSessionHelper.INSTANCE.get(mApp.getContext());
         mIconCache = mApp.getIconCache();
         mUserManagerState = userManagerState;
@@ -407,7 +409,6 @@
             @Nullable LoaderMemoryLogger memoryLogger,
             @Nullable LauncherRestoreEventLogger restoreEventLogger) {
         final Context context = mApp.getContext();
-        final PackageManagerHelper pmHelper = new PackageManagerHelper(context);
         final boolean isSdCardReady = Utilities.isBootCompleted();
         final WidgetInflater widgetInflater = new WidgetInflater(context);
 
@@ -447,7 +448,7 @@
                         mUserCache, mUserManagerState, mLauncherApps, mPendingPackages,
                         mShortcutKeyToPinnedShortcuts, mApp, mBgDataModel,
                         mWidgetProvidersMap, installingPkgs, isSdCardReady,
-                        widgetInflater, pmHelper, iconRequestInfos, unlockedUsers,
+                        widgetInflater, mPmHelper, iconRequestInfos, unlockedUsers,
                         allDeepShortcuts);
 
                 while (!mStopped && c.moveToNext()) {
diff --git a/src/com/android/launcher3/model/SdCardAvailableReceiver.java b/src/com/android/launcher3/model/SdCardAvailableReceiver.java
index 8cfa3aa..5293316 100644
--- a/src/com/android/launcher3/model/SdCardAvailableReceiver.java
+++ b/src/com/android/launcher3/model/SdCardAvailableReceiver.java
@@ -52,7 +52,7 @@
     @Override
     public void onReceive(Context context, Intent intent) {
         final LauncherApps launcherApps = context.getSystemService(LauncherApps.class);
-        final PackageManagerHelper pmHelper = new PackageManagerHelper(context);
+        final PackageManagerHelper pmHelper = PackageManagerHelper.INSTANCE.get(context);
         for (PackageUserKey puk : mPackages) {
             UserHandle user = puk.mUser;
 
diff --git a/src/com/android/launcher3/model/ShortcutsChangedTask.java b/src/com/android/launcher3/model/ShortcutsChangedTask.java
index 59dd1b1..1cb5215 100644
--- a/src/com/android/launcher3/model/ShortcutsChangedTask.java
+++ b/src/com/android/launcher3/model/ShortcutsChangedTask.java
@@ -78,7 +78,7 @@
 
         if (!matchingWorkspaceItems.isEmpty()) {
             if (mShortcuts.isEmpty()) {
-                PackageManagerHelper packageManagerHelper = new PackageManagerHelper(
+                PackageManagerHelper packageManagerHelper = PackageManagerHelper.INSTANCE.get(
                         app.getContext());
                 // Verify that the app is indeed installed.
                 if (!packageManagerHelper.isAppInstalled(mPackageName, mUser)
diff --git a/src/com/android/launcher3/model/data/ItemInfo.java b/src/com/android/launcher3/model/data/ItemInfo.java
index 8c3efd7..72e85c7 100644
--- a/src/com/android/launcher3/model/data/ItemInfo.java
+++ b/src/com/android/launcher3/model/data/ItemInfo.java
@@ -427,12 +427,10 @@
     @NonNull
     protected LauncherAtom.ItemInfo.Builder getDefaultItemInfoBuilder() {
         LauncherAtom.ItemInfo.Builder itemBuilder = LauncherAtom.ItemInfo.newBuilder();
-        UserIconInfo info = getUserInfo();
-        itemBuilder.setIsWork(info != null && info.isWork());
-        itemBuilder.setUserType(getUserType(info));
-        SettingsCache settingsCache = SettingsCache.INSTANCE.getNoCreate();
-        boolean isKidsMode = settingsCache != null && settingsCache.getValue(NAV_BAR_KIDS_MODE, 0);
-        itemBuilder.setIsKidsMode(isKidsMode);
+        SettingsCache.INSTANCE.executeIfCreated(cache ->
+                itemBuilder.setIsKidsMode(cache.getValue(NAV_BAR_KIDS_MODE, 0)));
+        UserCache.INSTANCE.executeIfCreated(cache ->
+                itemBuilder.setUserType(getUserType(cache.getUserInfo(user))));
         itemBuilder.setRank(rank);
         return itemBuilder;
     }
@@ -526,15 +524,6 @@
         this.title = title;
     }
 
-    private UserIconInfo getUserInfo() {
-        UserCache userCache = UserCache.INSTANCE.getNoCreate();
-        if (userCache == null) {
-            return null;
-        }
-
-        return userCache.getUserInfo(user);
-    }
-
     private int getUserType(UserIconInfo info) {
         if (info == null) {
             return SysUiStatsLog.LAUNCHER_UICHANGED__USER_TYPE__TYPE_UNKNOWN;
diff --git a/src/com/android/launcher3/pm/InstallSessionHelper.java b/src/com/android/launcher3/pm/InstallSessionHelper.java
index 4a3318e..e66f496 100644
--- a/src/com/android/launcher3/pm/InstallSessionHelper.java
+++ b/src/com/android/launcher3/pm/InstallSessionHelper.java
@@ -32,7 +32,6 @@
 import com.android.launcher3.Flags;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.SessionCommitReceiver;
-import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.ItemInstallQueue;
 import com.android.launcher3.util.IntArray;
@@ -41,6 +40,7 @@
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.SafeCloseable;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -52,7 +52,7 @@
  * Utility class to tracking install sessions
  */
 @SuppressWarnings("NewApi")
-public class InstallSessionHelper {
+public class InstallSessionHelper implements SafeCloseable {
 
     @NonNull
     private static final String LOG = "InstallSessionHelper";
@@ -89,6 +89,9 @@
         mLauncherApps = context.getSystemService(LauncherApps.class);
     }
 
+    @Override
+    public void close() { }
+
     @WorkerThread
     @NonNull
     private IntSet getPromiseIconIds() {
@@ -168,7 +171,7 @@
         synchronized (mSessionVerifiedMap) {
             if (!mSessionVerifiedMap.containsKey(pkg)) {
                 boolean hasSystemFlag = DEBUG || mAppContext.getPackageName().equals(pkg)
-                        || new PackageManagerHelper(mAppContext)
+                        || PackageManagerHelper.INSTANCE.get(mAppContext)
                                 .getApplicationInfo(pkg, user, ApplicationInfo.FLAG_SYSTEM) != null;
                 mSessionVerifiedMap.put(pkg, hasSystemFlag);
             }
@@ -242,7 +245,7 @@
                 && sessionInfo.getInstallReason() == PackageManager.INSTALL_REASON_USER
                 && sessionInfo.getAppIcon() != null
                 && !TextUtils.isEmpty(sessionInfo.getAppLabel())
-                && !new PackageManagerHelper(mAppContext).isAppInstalled(
+                && !PackageManagerHelper.INSTANCE.get(mAppContext).isAppInstalled(
                         sessionInfo.getAppPackageName(), getUserHandle(sessionInfo));
     }
 
diff --git a/src/com/android/launcher3/pm/UserCache.java b/src/com/android/launcher3/pm/UserCache.java
index 185e0b5..b7b557d 100644
--- a/src/com/android/launcher3/pm/UserCache.java
+++ b/src/com/android/launcher3/pm/UserCache.java
@@ -24,6 +24,7 @@
 import android.os.Process;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.util.ArrayMap;
 
 import androidx.annotation.AnyThread;
 import androidx.annotation.NonNull;
@@ -81,6 +82,9 @@
     @NonNull
     private Map<UserHandle, UserIconInfo> mUserToSerialMap;
 
+    @NonNull
+    private Map<UserHandle, List<String>> mUserToPreInstallAppMap;
+
     private UserCache(Context context) {
         mContext = context;
         mUserToSerialMap = Collections.emptyMap();
@@ -120,6 +124,20 @@
     @WorkerThread
     private void updateCache() {
         mUserToSerialMap = ApiWrapper.INSTANCE.get(mContext).queryAllUsers();
+        mUserToPreInstallAppMap = fetchPreInstallApps();
+    }
+
+    @WorkerThread
+    private Map<UserHandle, List<String>> fetchPreInstallApps() {
+        Map<UserHandle, List<String>> userToPreInstallApp = new ArrayMap<>();
+        mUserToSerialMap.forEach((userHandle, userIconInfo) -> {
+            // Fetch only for private profile, as other profiles have no usages yet.
+            List<String> preInstallApp = userIconInfo.isPrivate()
+                    ? ApiWrapper.INSTANCE.get(mContext).getPreInstalledSystemPackages(userHandle)
+                    : new ArrayList<>();
+            userToPreInstallApp.put(userHandle, preInstallApp);
+        });
+        return userToPreInstallApp;
     }
 
     /**
@@ -172,6 +190,15 @@
     }
 
     /**
+     * Returns the pre-installed apps for a user.
+     */
+    @NonNull
+    public List<String> getPreInstallApps(UserHandle user) {
+        List<String> preInstallApp = mUserToPreInstallAppMap.get(user);
+        return preInstallApp == null ? new ArrayList<>() : preInstallApp;
+    }
+
+    /**
      * Get a non-themed {@link UserBadgeDrawable} based on the provided {@link UserHandle}.
      */
     @Nullable
diff --git a/src/com/android/launcher3/popup/SystemShortcut.java b/src/com/android/launcher3/popup/SystemShortcut.java
index f56d732..6005573 100644
--- a/src/com/android/launcher3/popup/SystemShortcut.java
+++ b/src/com/android/launcher3/popup/SystemShortcut.java
@@ -181,8 +181,8 @@
         public void onClick(View view) {
             dismissTaskMenuView();
             Rect sourceBounds = Utilities.getViewBounds(view);
-            new PackageManagerHelper(view.getContext()).startDetailsActivityForInfo(
-                    mItemInfo, sourceBounds, ActivityOptions.makeBasic().toBundle());
+            PackageManagerHelper.startDetailsActivityForInfo(view.getContext(), mItemInfo,
+                    sourceBounds, ActivityOptions.makeBasic().toBundle());
             mTarget.getStatsLogManager().logger().withItemInfo(mItemInfo)
                     .log(LAUNCHER_SYSTEM_SHORTCUT_APP_INFO_TAP);
         }
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index d5f1e18..a4ff29f 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -526,10 +526,7 @@
         }
 
         logFavoritesTable(controller.getDb(), "launcher db after remap widget ids", null, null);
-        LauncherAppState app = LauncherAppState.getInstanceNoCreate();
-        if (app != null) {
-            app.getModel().forceReload();
-        }
+        LauncherAppState.INSTANCE.executeIfCreated(app -> app.getModel().forceReload());
     }
 
     private static void logDatabaseWidgetInfo(ModelDbController controller) {
diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index d5c9b9f..db2a6e0 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -94,13 +94,10 @@
 
     protected Context mContext;
     protected DeviceProfile mDeviceProfile;
-    protected LauncherAppState mLauncherAppState;
 
     public void init(Context context) {
         mContext = context;
-        mDeviceProfile = InvariantDeviceProfile.INSTANCE.
-                get(context).getDeviceProfile(context);
-        mLauncherAppState = LauncherAppState.getInstanceNoCreate();
+        mDeviceProfile = InvariantDeviceProfile.INSTANCE.get(context).getDeviceProfile(context);
         if (sActivityLifecycleCallbacks == null) {
             sActivityLifecycleCallbacks = new ActivityLifecycleCallbacksAdapter() {
                 @Override
diff --git a/src/com/android/launcher3/util/BgObjectWithLooper.java b/src/com/android/launcher3/util/BgObjectWithLooper.java
deleted file mode 100644
index adc3c7d..0000000
--- a/src/com/android/launcher3/util/BgObjectWithLooper.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * Copyright (C) 2020 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.database.ContentObserver;
-import android.net.Uri;
-import android.os.Handler;
-import android.os.Looper;
-
-import androidx.annotation.WorkerThread;
-
-import java.util.function.Consumer;
-
-/**
- * Utility class to define an object which does most of it's processing on a
- * dedicated background thread.
- */
-public abstract class BgObjectWithLooper {
-
-    /**
-     * Start initialization of the object
-     */
-    public final void initializeInBackground(String threadName) {
-        new Thread(this::runOnThread, threadName).start();
-    }
-
-    private void runOnThread() {
-        Looper.prepare();
-        onInitialized(Looper.myLooper());
-        Looper.loop();
-    }
-
-    /**
-     * Called on the background thread to handle initialization
-     */
-    @WorkerThread
-    protected abstract void onInitialized(Looper looper);
-
-    /**
-     * Helper method to create a content provider
-     */
-    protected static ContentObserver newContentObserver(Handler handler, Consumer<Uri> command) {
-        return new ContentObserver(handler) {
-            @Override
-            public void onChange(boolean selfChange, Uri uri) {
-                command.accept(uri);
-            }
-        };
-    }
-}
diff --git a/src/com/android/launcher3/util/DynamicResource.java b/src/com/android/launcher3/util/DynamicResource.java
index 1008ebb..fbdb5c2 100644
--- a/src/com/android/launcher3/util/DynamicResource.java
+++ b/src/com/android/launcher3/util/DynamicResource.java
@@ -33,7 +33,8 @@
  *
  * To allow customization for a particular resource, add them to dynamic_resources.xml
  */
-public class DynamicResource implements ResourceProvider, PluginListener<ResourceProvider> {
+public class DynamicResource implements
+        ResourceProvider, PluginListener<ResourceProvider>, SafeCloseable {
 
     private static final MainThreadInitializedObject<DynamicResource> INSTANCE =
             new MainThreadInitializedObject<>(DynamicResource::new);
@@ -48,6 +49,11 @@
     }
 
     @Override
+    public void close() {
+        PluginManagerWrapper.INSTANCE.get(mContext).removePluginListener(this);
+    }
+
+    @Override
     public int getInt(@IntegerRes int resId) {
         return mContext.getResources().getInteger(resId);
     }
diff --git a/src/com/android/launcher3/util/MainThreadInitializedObject.java b/src/com/android/launcher3/util/MainThreadInitializedObject.java
index b966d8e..1a0f9a0 100644
--- a/src/com/android/launcher3/util/MainThreadInitializedObject.java
+++ b/src/com/android/launcher3/util/MainThreadInitializedObject.java
@@ -28,17 +28,15 @@
 import com.android.launcher3.util.ResourceBasedOverride.Overrides;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.function.Consumer;
 
 /**
  * Utility class for defining singletons which are initiated on main thread.
  */
-public class MainThreadInitializedObject<T> {
+public class MainThreadInitializedObject<T extends SafeCloseable> {
 
     private final ObjectProvider<T> mProvider;
     private T mValue;
@@ -48,14 +46,14 @@
     }
 
     public T get(Context context) {
-        if (context instanceof SandboxContext sc) {
+        Context app = context.getApplicationContext();
+        if (app instanceof SandboxApplication sc) {
             return sc.getObject(this);
         }
 
         if (mValue == null) {
             if (Looper.myLooper() == Looper.getMainLooper()) {
-                mValue = TraceHelper.allowIpcs("main.thread.object",
-                        () -> mProvider.get(context.getApplicationContext()));
+                mValue = TraceHelper.allowIpcs("main.thread.object", () -> mProvider.get(app));
             } else {
                 try {
                     return MAIN_EXECUTOR.submit(() -> get(context)).get();
@@ -67,8 +65,18 @@
         return mValue;
     }
 
-    public T getNoCreate() {
-        return mValue;
+    /**
+     * Executes the callback is the value is already created
+     * @return true if the callback was executed, false otherwise
+     */
+    public boolean executeIfCreated(Consumer<T> callback) {
+        T v = mValue;
+        if (v != null) {
+            callback.accept(v);
+            return true;
+        } else {
+            return false;
+        }
     }
 
     @VisibleForTesting
@@ -79,8 +87,8 @@
     /**
      * Initializes a provider based on resource overrides
      */
-    public static <T extends ResourceBasedOverride> MainThreadInitializedObject<T> forOverride(
-            Class<T> clazz, int resourceId) {
+    public static <T extends ResourceBasedOverride & SafeCloseable> MainThreadInitializedObject<T>
+            forOverride(Class<T> clazz, int resourceId) {
         return new MainThreadInitializedObject<>(c -> Overrides.getObject(clazz, c, resourceId));
     }
 
@@ -89,24 +97,36 @@
         T get(Context context);
     }
 
+    public interface SandboxApplication {
+
+        /**
+         * Find a cached object from mObjectMap if we have already created one. If not, generate
+         * an object using the provider.
+         */
+        <T extends SafeCloseable> T getObject(MainThreadInitializedObject<T> object);
+
+        @UiThread
+        default <T extends SafeCloseable> T createObject(MainThreadInitializedObject<T> object) {
+            return object.mProvider.get((Context) this);
+        }
+    }
+
     /**
      * Abstract Context which allows custom implementations for
      * {@link MainThreadInitializedObject} providers
      */
-    public static class SandboxContext extends ContextWrapper {
+    public static class SandboxContext extends ContextWrapper implements SandboxApplication {
 
         private static final String TAG = "SandboxContext";
 
-        protected final Set<MainThreadInitializedObject> mAllowedObjects;
-        protected final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
-        protected final ArrayList<Object> mOrderedObjects = new ArrayList<>();
+        private final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
+        private final ArrayList<SafeCloseable> mOrderedObjects = new ArrayList<>();
 
         private final Object mDestroyLock = new Object();
         private boolean mDestroyed = false;
 
-        public SandboxContext(Context base, MainThreadInitializedObject... allowedObjects) {
+        public SandboxContext(Context base) {
             super(base);
-            mAllowedObjects = new HashSet<>(Arrays.asList(allowedObjects));
         }
 
         @Override
@@ -118,20 +138,14 @@
             synchronized (mDestroyLock) {
                 // Destroy in reverse order
                 for (int i = mOrderedObjects.size() - 1; i >= 0; i--) {
-                    Object o = mOrderedObjects.get(i);
-                    if (o instanceof SafeCloseable) {
-                        ((SafeCloseable) o).close();
-                    }
+                    mOrderedObjects.get(i).close();
                 }
                 mDestroyed = true;
             }
         }
 
-        /**
-         * Find a cached object from mObjectMap if we have already created one. If not, generate
-         * an object using the provider.
-         */
-        protected <T> T getObject(MainThreadInitializedObject<T> object) {
+        @Override
+        public <T extends SafeCloseable> T getObject(MainThreadInitializedObject<T> object) {
             synchronized (mDestroyLock) {
                 if (mDestroyed) {
                     Log.e(TAG, "Static object access with a destroyed context");
@@ -142,12 +156,6 @@
                 }
                 if (Looper.myLooper() == Looper.getMainLooper()) {
                     t = createObject(object);
-                    // Check if we've explicitly allowed the object or if it's a SafeCloseable,
-                    // it will get destroyed in onDestroy()
-                    if (!mAllowedObjects.contains(object) && !(t instanceof SafeCloseable)) {
-                        throw new IllegalStateException("Leaking unknown objects "
-                                + object + "  " + object.mProvider + " " + t);
-                    }
                     mObjectMap.put(object, t);
                     mOrderedObjects.add(t);
                     return t;
@@ -161,17 +169,12 @@
             }
         }
 
-        @UiThread
-        protected <T> T createObject(MainThreadInitializedObject<T> object) {
-            return object.mProvider.get(this);
-        }
-
         /**
          * Put a value into mObjectMap, can be used to put mocked MainThreadInitializedObject
          * instances into SandboxContext.
          */
-        @VisibleForTesting
-        public <T> void putObject(MainThreadInitializedObject<T> object, T value) {
+        public <T extends SafeCloseable> void putObject(
+                MainThreadInitializedObject<T> object, T value) {
             mObjectMap.put(object, value);
         }
     }
diff --git a/src/com/android/launcher3/util/PackageManagerHelper.java b/src/com/android/launcher3/util/PackageManagerHelper.java
index 608bed7..3684f56 100644
--- a/src/com/android/launcher3/util/PackageManagerHelper.java
+++ b/src/com/android/launcher3/util/PackageManagerHelper.java
@@ -56,11 +56,15 @@
 /**
  * Utility methods using package manager
  */
-public class PackageManagerHelper {
+public class PackageManagerHelper implements SafeCloseable{
 
     private static final String TAG = "PackageManagerHelper";
 
     @NonNull
+    public static final MainThreadInitializedObject<PackageManagerHelper> INSTANCE =
+            new MainThreadInitializedObject<>(PackageManagerHelper::new);
+
+    @NonNull
     private final Context mContext;
 
     @NonNull
@@ -75,6 +79,9 @@
         mLauncherApps = Objects.requireNonNull(context.getSystemService(LauncherApps.class));
     }
 
+    @Override
+    public void close() { }
+
     /**
      * Returns true if the app can possibly be on the SDCard. This is just a workaround and doesn't
      * guarantee that the app is on SD card.
@@ -170,10 +177,11 @@
     /**
      * Starts the details activity for {@code info}
      */
-    public void startDetailsActivityForInfo(ItemInfo info, Rect sourceBounds, Bundle opts) {
+    public static void startDetailsActivityForInfo(Context context, ItemInfo info,
+            Rect sourceBounds, Bundle opts) {
         if (info instanceof ItemInfoWithIcon appInfo
                 && (appInfo.runtimeStatusFlags & FLAG_INSTALL_SESSION_ACTIVE) != 0) {
-            mContext.startActivity(ApiWrapper.INSTANCE.get(mContext).getAppMarketActivityIntent(
+            context.startActivity(ApiWrapper.INSTANCE.get(context).getAppMarketActivityIntent(
                     appInfo.getTargetComponent().getPackageName(), Process.myUserHandle()));
             return;
         }
@@ -189,9 +197,10 @@
         }
         if (componentName != null) {
             try {
-                mLauncherApps.startAppDetailsActivity(componentName, info.user, sourceBounds, opts);
+                context.getSystemService(LauncherApps.class).startAppDetailsActivity(componentName,
+                        info.user, sourceBounds, opts);
             } catch (SecurityException | ActivityNotFoundException e) {
-                Toast.makeText(mContext, R.string.activity_not_found, Toast.LENGTH_SHORT).show();
+                Toast.makeText(context, R.string.activity_not_found, Toast.LENGTH_SHORT).show();
                 Log.e(TAG, "Unable to launch settings", e);
             }
         }
diff --git a/src/com/android/launcher3/util/ScreenOnTracker.java b/src/com/android/launcher3/util/ScreenOnTracker.java
index 67530a6..e16e477 100644
--- a/src/com/android/launcher3/util/ScreenOnTracker.java
+++ b/src/com/android/launcher3/util/ScreenOnTracker.java
@@ -27,7 +27,7 @@
 /**
  * Utility class for tracking if the screen is currently on or off
  */
-public class ScreenOnTracker {
+public class ScreenOnTracker implements SafeCloseable {
 
     public static final MainThreadInitializedObject<ScreenOnTracker> INSTANCE =
             new MainThreadInitializedObject<>(ScreenOnTracker::new);
@@ -35,14 +35,21 @@
     private final SimpleBroadcastReceiver mReceiver = new SimpleBroadcastReceiver(this::onReceive);
     private final CopyOnWriteArrayList<ScreenOnListener> mListeners = new CopyOnWriteArrayList<>();
 
+    private final Context mContext;
     private boolean mIsScreenOn;
 
     private ScreenOnTracker(Context context) {
         // Assume that the screen is on to begin with
+        mContext = context;
         mIsScreenOn = true;
         mReceiver.register(context, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT);
     }
 
+    @Override
+    public void close() {
+        mReceiver.unregisterReceiverSafely(mContext);
+    }
+
     private void onReceive(Intent intent) {
         String action = intent.getAction();
         if (ACTION_SCREEN_ON.equals(action)) {
diff --git a/src/com/android/launcher3/util/VibratorWrapper.java b/src/com/android/launcher3/util/VibratorWrapper.java
index cd60c1d..51749a7 100644
--- a/src/com/android/launcher3/util/VibratorWrapper.java
+++ b/src/com/android/launcher3/util/VibratorWrapper.java
@@ -19,14 +19,12 @@
 import static android.os.VibrationEffect.createPredefined;
 import static android.provider.Settings.System.HAPTIC_FEEDBACK_ENABLED;
 
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 
 import android.annotation.SuppressLint;
-import android.content.ContentResolver;
 import android.content.Context;
-import android.database.ContentObserver;
 import android.media.AudioAttributes;
+import android.net.Uri;
 import android.os.SystemClock;
 import android.os.VibrationEffect;
 import android.os.Vibrator;
@@ -39,7 +37,7 @@
 /**
  * Wrapper around {@link Vibrator} to easily perform haptic feedback where necessary.
  */
-public class VibratorWrapper {
+public class VibratorWrapper implements SafeCloseable {
 
     public static final MainThreadInitializedObject<VibratorWrapper> INSTANCE =
             new MainThreadInitializedObject<>(VibratorWrapper::new);
@@ -51,6 +49,8 @@
 
     public static final VibrationEffect EFFECT_CLICK =
             createPredefined(VibrationEffect.EFFECT_CLICK);
+    private static final Uri HAPTIC_FEEDBACK_URI =
+            Settings.System.getUriFor(HAPTIC_FEEDBACK_ENABLED);
 
     private static final float LOW_TICK_SCALE = 0.9f;
     private static final float DRAG_TEXTURE_SCALE = 0.03f;
@@ -76,6 +76,8 @@
     private final Context mContext;
     private final Vibrator mVibrator;
     private final boolean mHasVibrator;
+    private final SettingsCache.OnChangeListener mHapticChangeListener =
+            isEnabled -> mIsHapticFeedbackEnabled = isEnabled;
 
     private boolean mIsHapticFeedbackEnabled;
 
@@ -84,16 +86,9 @@
         mVibrator = context.getSystemService(Vibrator.class);
         mHasVibrator = mVibrator.hasVibrator();
         if (mHasVibrator) {
-            final ContentResolver resolver = context.getContentResolver();
-            mIsHapticFeedbackEnabled = isHapticFeedbackEnabled(resolver);
-            final ContentObserver observer = new ContentObserver(MAIN_EXECUTOR.getHandler()) {
-                @Override
-                public void onChange(boolean selfChange) {
-                    mIsHapticFeedbackEnabled = isHapticFeedbackEnabled(resolver);
-                }
-            };
-            resolver.registerContentObserver(Settings.System.getUriFor(HAPTIC_FEEDBACK_ENABLED),
-                    false /* notifyForDescendants */, observer);
+            SettingsCache cache = SettingsCache.INSTANCE.get(mContext);
+            cache.register(HAPTIC_FEEDBACK_URI, mHapticChangeListener);
+            mIsHapticFeedbackEnabled = cache.getValue(HAPTIC_FEEDBACK_URI, 0);
         } else {
             mIsHapticFeedbackEnabled = false;
         }
@@ -126,6 +121,14 @@
         }
     }
 
+    @Override
+    public void close() {
+        if (mHasVibrator) {
+            SettingsCache.INSTANCE.get(mContext)
+                    .unregister(HAPTIC_FEEDBACK_URI, mHapticChangeListener);
+        }
+    }
+
     /**
      * This is called when the user swipes to/from all apps. This is meant to be used in between
      * long animation progresses so that it gives a dragging texture effect. For a better
@@ -175,10 +178,6 @@
         mLastDragTime = 0;
     }
 
-    private boolean isHapticFeedbackEnabled(ContentResolver resolver) {
-        return Settings.System.getInt(resolver, HAPTIC_FEEDBACK_ENABLED, 0) == 1;
-    }
-
     /** Vibrates with the given effect if haptic feedback is available and enabled. */
     public void vibrate(VibrationEffect vibrationEffect) {
         if (mHasVibrator && mIsHapticFeedbackEnabled) {
diff --git a/src/com/android/launcher3/util/window/WindowManagerProxy.java b/src/com/android/launcher3/util/window/WindowManagerProxy.java
index 998191e..4b004f3 100644
--- a/src/com/android/launcher3/util/window/WindowManagerProxy.java
+++ b/src/com/android/launcher3/util/window/WindowManagerProxy.java
@@ -59,6 +59,7 @@
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.NavigationMode;
 import com.android.launcher3.util.ResourceBasedOverride;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.WindowBounds;
 
 import java.util.ArrayList;
@@ -67,7 +68,7 @@
 /**
  * Utility class for mocking some window manager behaviours
  */
-public class WindowManagerProxy implements ResourceBasedOverride {
+public class WindowManagerProxy implements ResourceBasedOverride, SafeCloseable {
 
     private static final String TAG = "WindowManagerProxy";
     public static final int MIN_TABLET_WIDTH = 600;
@@ -305,12 +306,12 @@
 
         navBarHeightPortrait = isTablet
                 ? (mTaskbarDrawnInProcess
-                        ? 0 : systemRes.getDimensionPixelSize(R.dimen.taskbar_size))
+                        ? 0 : context.getResources().getDimensionPixelSize(R.dimen.taskbar_size))
                 : getDimenByName(systemRes, NAVBAR_HEIGHT);
 
         navBarHeightLandscape = isTablet
                 ? (mTaskbarDrawnInProcess
-                        ? 0 : systemRes.getDimensionPixelSize(R.dimen.taskbar_size))
+                        ? 0 : context.getResources().getDimensionPixelSize(R.dimen.taskbar_size))
                 : (isTabletOrGesture
                         ? getDimenByName(systemRes, NAVBAR_HEIGHT_LANDSCAPE) : 0);
         navbarWidthLandscape = isTabletOrGesture
@@ -474,6 +475,9 @@
                 NavigationMode.THREE_BUTTONS;
     }
 
+    @Override
+    public void close() { }
+
     /**
      * @see DisplayCutout#getSafeInsets
      */
diff --git a/src/com/android/launcher3/widget/BaseWidgetSheet.java b/src/com/android/launcher3/widget/BaseWidgetSheet.java
index e5b5daa..749554c 100644
--- a/src/com/android/launcher3/widget/BaseWidgetSheet.java
+++ b/src/com/android/launcher3/widget/BaseWidgetSheet.java
@@ -141,9 +141,17 @@
 
         if (enableWidgetTapToAdd()) {
             scrollToWidgetCell(wc);
+
             if (mWidgetCellWithAddButton != null) {
-                // If there is a add button currently showing, hide it.
-                mWidgetCellWithAddButton.hideAddButton(/* animate= */ true);
+                if (mWidgetCellWithAddButton.isShowingAddButton()) {
+                    // If there is a add button currently showing, hide it.
+                    mWidgetCellWithAddButton.hideAddButton(/* animate= */ true);
+                } else {
+                    // The last recorded widget cell to show an add button is no longer showing it,
+                    // likely because the widget cell has been recycled or lost focus. If this is
+                    // the cell that has been clicked, we will show it below.
+                    mWidgetCellWithAddButton = null;
+                }
             }
 
             if (mWidgetCellWithAddButton != wc) {
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index ab007fb..ea167d7 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -110,6 +110,7 @@
     private int mSourceContainer = CONTAINER_WIDGETS_TRAY;
 
     private CancellableTask mIconLoadRequest;
+    private boolean mIsShowingAddButton = false;
 
     public WidgetCell(Context context) {
         this(context, null);
@@ -534,6 +535,9 @@
      * @param callback Callback to be set on the button.
      */
     public void showAddButton(View.OnClickListener callback) {
+        if (mIsShowingAddButton) return;
+        mIsShowingAddButton = true;
+
         mWidgetAddButton.setAlpha(0F);
         mWidgetAddButton.setVisibility(VISIBLE);
         mWidgetAddButton.setOnClickListener(callback);
@@ -555,6 +559,9 @@
      * Hide tap to add button.
      */
     public void hideAddButton(boolean animate) {
+        if (!mIsShowingAddButton) return;
+        mIsShowingAddButton = false;
+
         mWidgetAddButton.setOnClickListener(null);
         mWidgetAddButton.animate().cancel();
         mWidgetTextContainer.animate().cancel();
@@ -579,4 +586,8 @@
                 .alpha(1F)
                 .setDuration(ADD_BUTTON_FADE_DURATION_MS);
     }
+
+    public boolean isShowingAddButton() {
+        return mIsShowingAddButton;
+    }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
index c3906a6..80bda22 100644
--- a/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
+++ b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
@@ -62,16 +62,17 @@
         // via the overridden WidgetRecommendationCategoryProvider resource.
 
         Preconditions.assertWorkerThread();
-        PackageManagerHelper pmHelper = new PackageManagerHelper(context);
-        if (item.widgetInfo != null && item.widgetInfo.getComponent() != null) {
-            String widgetComponentName = item.widgetInfo.getComponent().getClassName();
-            ApplicationInfo applicationInfo = pmHelper.getApplicationInfo(
-                    item.widgetInfo.getComponent().getPackageName(), item.widgetInfo.getUser(),
-                    0 /* flags */);
-            if (applicationInfo != null) {
-                int predictionCategory = applicationInfo.category;
-                return getCategoryFromApplicationCategory(context, predictionCategory,
-                        widgetComponentName);
+        try (PackageManagerHelper pmHelper = new PackageManagerHelper(context)) {
+            if (item.widgetInfo != null && item.widgetInfo.getComponent() != null) {
+                String widgetComponentName = item.widgetInfo.getComponent().getClassName();
+                ApplicationInfo applicationInfo = pmHelper.getApplicationInfo(
+                        item.widgetInfo.getComponent().getPackageName(), item.widgetInfo.getUser(),
+                        0 /* flags */);
+                if (applicationInfo != null) {
+                    int predictionCategory = applicationInfo.category;
+                    return getCategoryFromApplicationCategory(context, predictionCategory,
+                            widgetComponentName);
+                }
             }
         }
         return null;
diff --git a/tests/OWNERS b/tests/OWNERS
index b5ee7d7..6b8643c 100644
--- a/tests/OWNERS
+++ b/tests/OWNERS
@@ -3,4 +3,3 @@
 sunnygoyal@google.com
 winsonc@google.com
 hyunyoungs@google.com
-mateuszc@google.com
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
index 0009a4e..002f496 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -52,17 +52,11 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherModel.ModelUpdateTask;
-import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.model.AllAppsList;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.ItemInstallQueue;
-import com.android.launcher3.pm.InstallSessionHelper;
-import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.testing.TestInformationProvider;
 import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
-import com.android.launcher3.util.window.WindowManagerProxy;
-import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.io.ByteArrayOutputStream;
 import java.io.File;
@@ -233,13 +227,7 @@
         private final File mDbDir;
 
         public SandboxModelContext() {
-            super(ApplicationProvider.getApplicationContext(),
-                    UserCache.INSTANCE, InstallSessionHelper.INSTANCE, LauncherPrefs.INSTANCE,
-                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
-                    DisplayController.INSTANCE, CustomWidgetManager.INSTANCE,
-                    SettingsCache.INSTANCE, PluginManagerWrapper.INSTANCE,
-                    LockedUserState.INSTANCE, WallpaperColorHints.INSTANCE,
-                    ItemInstallQueue.INSTANCE, WindowManagerProxy.INSTANCE);
+            super(ApplicationProvider.getApplicationContext());
 
             // System settings cache content provider. Ensure that they are statically initialized
             Settings.Secure.getString(
@@ -254,18 +242,13 @@
         }
 
         @Override
-        protected <T> T createObject(MainThreadInitializedObject<T> object) {
+        public <T extends SafeCloseable> T createObject(MainThreadInitializedObject<T> object) {
             if (object == LauncherAppState.INSTANCE) {
                 return (T) new LauncherAppState(this, null /* iconCacheFileName */);
             }
             return super.createObject(object);
         }
 
-        public SandboxModelContext allow(MainThreadInitializedObject object) {
-            mAllowedObjects.add(object);
-            return this;
-        }
-
         @Override
         public File getDatabasePath(String name) {
             if (!mDbDir.exists()) {
diff --git a/tests/shared/com/android/launcher3/testing/OWNERS b/tests/shared/com/android/launcher3/testing/OWNERS
new file mode 100644
index 0000000..02e8ebc
--- /dev/null
+++ b/tests/shared/com/android/launcher3/testing/OWNERS
@@ -0,0 +1,4 @@
+vadimt@google.com
+sunnygoyal@google.com
+winsonc@google.com
+hyunyoungs@google.com
diff --git a/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt b/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
index 577334b..aefc2db 100644
--- a/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
+++ b/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
@@ -300,13 +300,7 @@
                 smallestScreenWidthDp = min(screenWidthDp, screenHeightDp)
             }
         val configurationContext = runningContext.createConfigurationContext(config)
-        context =
-            SandboxContext(
-                configurationContext,
-                DisplayController.INSTANCE,
-                WindowManagerProxy.INSTANCE,
-                LauncherPrefs.INSTANCE
-            )
+        context = SandboxContext(configurationContext)
         context.putObject(DisplayController.INSTANCE, displayController)
         context.putObject(WindowManagerProxy.INSTANCE, windowManagerProxy)
         context.putObject(LauncherPrefs.INSTANCE, launcherPrefs)
diff --git a/tests/src/com/android/launcher3/allapps/AlphabeticalAppsListTest.java b/tests/src/com/android/launcher3/allapps/AlphabeticalAppsListTest.java
index 423ca24..d2238ff 100644
--- a/tests/src/com/android/launcher3/allapps/AlphabeticalAppsListTest.java
+++ b/tests/src/com/android/launcher3/allapps/AlphabeticalAppsListTest.java
@@ -98,7 +98,7 @@
         when(mPrivateProfileManager.addPrivateSpaceHeader(any()))
                 .thenAnswer(answer(this::addPrivateSpaceHeader));
         when(mPrivateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps())
+        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps(any()))
                 .thenReturn(iteminfo -> iteminfo.componentName == null
                         || !iteminfo.componentName.getPackageName()
                         .equals("com.android.launcher3.tests.camera"));
@@ -127,7 +127,7 @@
         when(mPrivateProfileManager.addSystemAppsDivider(any()))
                 .thenAnswer(answer(this::addSystemAppsDivider));
         when(mPrivateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps())
+        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps(mContext))
                 .thenReturn(iteminfo -> iteminfo.componentName == null
                         || !iteminfo.componentName.getPackageName()
                         .equals("com.android.launcher3.tests.camera"));
diff --git a/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewTest.java b/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewTest.java
index 9363ce8..512b2ac 100644
--- a/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewTest.java
+++ b/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewTest.java
@@ -71,6 +71,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Predicate;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -279,10 +280,8 @@
         when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
         PrivateProfileManager privateProfileManager = spy(mPrivateProfileManager);
         when(privateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(privateProfileManager.splitIntoUserInstalledAndSystemApps())
-                .thenReturn(iteminfo -> iteminfo.componentName == null
-                        || !iteminfo.componentName.getPackageName()
-                        .equals(CAMERA_PACKAGE_NAME));
+        doReturn(splitIntoUserInstalledAndSystemApps()).when(privateProfileManager)
+                .splitIntoUserInstalledAndSystemApps(any());
         doReturn(0).when(privateProfileManager).addPrivateSpaceHeader(any());
         doAnswer(answer(this::addPrivateSpaceHeader)).when(privateProfileManager)
                 .addPrivateSpaceHeader(any());
@@ -316,10 +315,8 @@
         when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
         PrivateProfileManager privateProfileManager = spy(mPrivateProfileManager);
         when(privateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(privateProfileManager.splitIntoUserInstalledAndSystemApps())
-                .thenReturn(iteminfo -> iteminfo.componentName == null
-                        || !iteminfo.componentName.getPackageName()
-                        .equals(CAMERA_PACKAGE_NAME));
+        doReturn(splitIntoUserInstalledAndSystemApps()).when(privateProfileManager)
+                .splitIntoUserInstalledAndSystemApps(any());
         doReturn(0).when(privateProfileManager).addPrivateSpaceHeader(any());
         doAnswer(answer(this::addPrivateSpaceHeader)).when(privateProfileManager)
                 .addPrivateSpaceHeader(any());
@@ -353,10 +350,8 @@
         when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
         PrivateProfileManager privateProfileManager = spy(mPrivateProfileManager);
         when(privateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(privateProfileManager.splitIntoUserInstalledAndSystemApps())
-                .thenReturn(iteminfo -> iteminfo.componentName == null
-                        || !iteminfo.componentName.getPackageName()
-                        .equals(CAMERA_PACKAGE_NAME));
+        doReturn(splitIntoUserInstalledAndSystemApps()).when(privateProfileManager)
+                .splitIntoUserInstalledAndSystemApps(any());
         doReturn(0).when(privateProfileManager).addPrivateSpaceHeader(any());
         doAnswer(answer(this::addPrivateSpaceHeader)).when(privateProfileManager)
                 .addPrivateSpaceHeader(any());
@@ -390,10 +385,8 @@
         when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
         PrivateProfileManager privateProfileManager = spy(mPrivateProfileManager);
         when(privateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(privateProfileManager.splitIntoUserInstalledAndSystemApps())
-                .thenReturn(iteminfo -> iteminfo.componentName == null
-                        || !iteminfo.componentName.getPackageName()
-                        .equals(CAMERA_PACKAGE_NAME));
+        doReturn(splitIntoUserInstalledAndSystemApps()).when(privateProfileManager)
+                .splitIntoUserInstalledAndSystemApps(any());
         doReturn(0).when(privateProfileManager).addPrivateSpaceHeader(any());
         doNothing().when(privateProfileManager).addPrivateSpaceInstallAppButton(any());
         doReturn(0).when(privateProfileManager).addSystemAppsDivider(any());
@@ -465,4 +458,10 @@
         }
         return appInfos.toArray(AppInfo[]::new);
     }
+
+    private Predicate<AppInfo> splitIntoUserInstalledAndSystemApps() {
+        return iteminfo -> iteminfo.componentName == null
+                || !iteminfo.componentName.getPackageName()
+                .equals(CAMERA_PACKAGE_NAME);
+    }
 }
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
index 342eedf..4a67600 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
@@ -21,8 +21,9 @@
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.Launcher;
+import com.android.launcher3.util.rule.ScreenRecordRule.KeepRecordOnSuccess;
+import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -31,6 +32,7 @@
 @RunWith(AndroidJUnit4.class)
 public class TaplTestsLauncher3Test extends AbstractLauncherUiTest<Launcher> {
 
+    @KeepRecordOnSuccess
     @ScreenRecord // b/322823478
     @Test
     public void testDevicePressMenu() throws Exception {
diff --git a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java b/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
index e08d37c..70a3dd0 100644
--- a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
@@ -45,6 +45,7 @@
 import com.android.launcher3.allapps.WorkProfileManager;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.util.TestUtil;
+import com.android.launcher3.util.rule.ScreenRecordRule.KeepRecordOnSuccess;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.util.rule.TestStabilityRule.Stability;
 
@@ -196,6 +197,7 @@
 
     }
 
+    @KeepRecordOnSuccess
     @ScreenRecord // b/322823478
     @Test
     public void testEdu() {
diff --git a/tests/src/com/android/launcher3/util/DisplayControllerTest.kt b/tests/src/com/android/launcher3/util/DisplayControllerTest.kt
index 706ab27..2e57ad5 100644
--- a/tests/src/com/android/launcher3/util/DisplayControllerTest.kt
+++ b/tests/src/com/android/launcher3/util/DisplayControllerTest.kt
@@ -45,6 +45,7 @@
 import org.mockito.kotlin.any
 import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.doNothing
+import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
@@ -123,6 +124,7 @@
         whenever(displayManager.getDisplay(any())).thenReturn(display)
 
         // Mock resources
+        doReturn(context).whenever(context).applicationContext
         whenever(resources.configuration).thenReturn(configuration)
         whenever(context.resources).thenReturn(resources)
 
diff --git a/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java b/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
index 7a5cf2c..ff96afb 100644
--- a/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
+++ b/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
@@ -49,6 +49,9 @@
             return base;
         }
 
+        final Boolean keepRecordOnSuccess = description.getAnnotation(KeepRecordOnSuccess.class)
+                != null;
+
         return new Statement() {
             @Override
             public void evaluate() throws Throwable {
@@ -70,7 +73,7 @@
                     device.executeShellCommand("kill -INT " + screenRecordPid);
                     Log.e(TAG, "Screenrecord captured at: " + outputFile);
                     output.close();
-                    if (success) {
+                    if (success && !keepRecordOnSuccess) {
                         automation.executeShellCommand("rm " + outputFile);
                     }
                 }
@@ -85,4 +88,14 @@
     @Target(ElementType.METHOD)
     public @interface ScreenRecord {
     }
+
+
+    /**
+     * Interface to indicate that we should keep the screen record even if the test succeeds, use
+     * sparingly since it uses a lof of memory.
+     */
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.METHOD)
+    public @interface KeepRecordOnSuccess {
+    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java
index 2c16e01..daddd2f 100644
--- a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java
+++ b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.tapl;
 
+import android.graphics.Point;
+
 import androidx.test.uiautomator.UiObject2;
 
 /**
@@ -57,4 +59,16 @@
     private void verifyDividerIsPresent() {
         mLauncher.waitForObjectInContainer(mAppListRecycler, DIVIDER_RES_ID);
     }
+
+    /**
+     * Verifies that a user installed app is present above the divider.
+     */
+    public void verifyInstalledAppIsPresent(String appName) {
+        HomeAppIcon appIcon = mLauncher.getAllApps().getAppIcon(appName);
+        final Point iconCenter = appIcon.mObject.getVisibleCenter();
+        UiObject2 divider = mLauncher.waitForObjectInContainer(mAppListRecycler, DIVIDER_RES_ID);
+        final Point dividerCenter = divider.getVisibleCenter();
+        mLauncher.assertTrue("Installed App: " + appName + " is not above the divider",
+                iconCenter.y < dividerCenter.y);
+    }
 }