diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index 457fdd8..aafa1f6 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -448,6 +448,12 @@
     bug: "292269949"
 }
 
+flag {
+    name: "enable_state_manager_proto_log"
+    namespace: "launcher"
+    description: "Enables tracking state manager logs in ProtoLog"
+    bug: "292269949"
+}
 
 flag {
     name: "coordinate_workspace_scale"
diff --git a/quickstep/res/drawable/ic_external_display.xml b/quickstep/res/drawable/ic_external_display.xml
new file mode 100644
index 0000000..64c183e
--- /dev/null
+++ b/quickstep/res/drawable/ic_external_display.xml
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2023 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960">
+    <path
+        android:pathData="M320,840v-80L160,760q-33,0 -56.5,-23.5T80,680v-480q0,-33 23.5,-56.5T160,120h640q33,0 56.5,23.5T880,200v480q0,33 -23.5,56.5T800,760L640,760v80L320,840ZM160,680h640v-480L160,200v480ZM160,680v-480,480Z"
+        android:fillColor="#e8eaed"/>
+</vector>
diff --git a/quickstep/res/layout/customizable_taskbar.xml b/quickstep/res/layout/customizable_taskbar.xml
index e1a80ae..d988cbc 100644
--- a/quickstep/res/layout/customizable_taskbar.xml
+++ b/quickstep/res/layout/customizable_taskbar.xml
@@ -51,20 +51,26 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent"/>
 
-    <com.android.launcher3.taskbar.bubbles.BubbleBarView
-        android:id="@+id/taskbar_bubbles"
-        android:layout_width="wrap_content"
-        android:layout_height="@dimen/bubblebar_size_with_pointer"
-        android:layout_gravity="bottom|end"
-        android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
-        android:paddingTop="@dimen/bubblebar_pointer_visible_size"
-        android:paddingEnd="@dimen/taskbar_icon_spacing"
-        android:paddingStart="@dimen/taskbar_icon_spacing"
-        android:visibility="gone"
-        android:gravity="center"
-        android:clipChildren="false"
-        android:elevation="@dimen/bubblebar_elevation"
-        />
+    <FrameLayout
+        android:id="@+id/taskbar_bubbles_container"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:clipChildren="false">
+
+        <com.android.launcher3.taskbar.bubbles.BubbleBarView
+            android:id="@+id/taskbar_bubbles"
+            android:layout_width="wrap_content"
+            android:layout_height="@dimen/bubblebar_size_with_pointer"
+            android:layout_gravity="bottom|end"
+            android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
+            android:paddingTop="@dimen/bubblebar_pointer_visible_size"
+            android:paddingEnd="@dimen/taskbar_icon_spacing"
+            android:paddingStart="@dimen/taskbar_icon_spacing"
+            android:visibility="gone"
+            android:gravity="center"
+            android:clipChildren="false"
+            android:elevation="@dimen/bubblebar_elevation" />
+    </FrameLayout>
 
     <com.android.launcher3.taskbar.navbutton.NearestTouchFrame
         android:id="@+id/navbuttons_view"
diff --git a/quickstep/res/layout/taskbar.xml b/quickstep/res/layout/taskbar.xml
index e8f3d9d..54f9ae8 100644
--- a/quickstep/res/layout/taskbar.xml
+++ b/quickstep/res/layout/taskbar.xml
@@ -35,17 +35,24 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent"/>
 
-    <com.android.launcher3.taskbar.bubbles.BubbleBarView
-        android:id="@+id/taskbar_bubbles"
-        android:layout_width="wrap_content"
-        android:layout_height="@dimen/bubblebar_size_with_pointer"
-        android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
-        android:paddingTop="@dimen/bubblebar_pointer_visible_size"
-        android:visibility="gone"
-        android:gravity="center"
-        android:layout_gravity="bottom"
-        android:clipChildren="false"
-        android:elevation="@dimen/bubblebar_elevation" />
+    <FrameLayout
+        android:id="@+id/taskbar_bubbles_container"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:clipChildren="false">
+
+        <com.android.launcher3.taskbar.bubbles.BubbleBarView
+            android:id="@+id/taskbar_bubbles"
+            android:layout_width="wrap_content"
+            android:layout_height="@dimen/bubblebar_size_with_pointer"
+            android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
+            android:paddingTop="@dimen/bubblebar_pointer_visible_size"
+            android:visibility="gone"
+            android:gravity="center"
+            android:layout_gravity="bottom"
+            android:clipChildren="false"
+            android:elevation="@dimen/bubblebar_elevation" />
+    </FrameLayout>
 
     <com.android.launcher3.taskbar.navbutton.NearestTouchFrame
         android:id="@+id/navbuttons_view"
diff --git a/quickstep/res/layout/transient_taskbar.xml b/quickstep/res/layout/transient_taskbar.xml
index f3c3383..3ec8046 100644
--- a/quickstep/res/layout/transient_taskbar.xml
+++ b/quickstep/res/layout/transient_taskbar.xml
@@ -38,18 +38,24 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent"/>
 
-    <com.android.launcher3.taskbar.bubbles.BubbleBarView
-        android:id="@+id/taskbar_bubbles"
-        android:layout_width="wrap_content"
-        android:layout_height="@dimen/bubblebar_size_with_pointer"
-        android:layout_gravity="bottom|end"
-        android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
-        android:paddingTop="@dimen/bubblebar_pointer_visible_size"
-        android:visibility="gone"
-        android:gravity="center"
-        android:clipChildren="false"
-        android:elevation="@dimen/bubblebar_elevation"
-        />
+    <FrameLayout
+        android:id="@+id/taskbar_bubbles_container"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:clipChildren="false">
+
+        <com.android.launcher3.taskbar.bubbles.BubbleBarView
+            android:id="@+id/taskbar_bubbles"
+            android:layout_width="wrap_content"
+            android:layout_height="@dimen/bubblebar_size_with_pointer"
+            android:layout_gravity="bottom|end"
+            android:layout_marginHorizontal="@dimen/transient_taskbar_bottom_margin"
+            android:paddingTop="@dimen/bubblebar_pointer_visible_size"
+            android:visibility="gone"
+            android:gravity="center"
+            android:clipChildren="false"
+            android:elevation="@dimen/bubblebar_elevation" />
+    </FrameLayout>
 
     <com.android.launcher3.taskbar.navbutton.NearestTouchFrame
         android:id="@+id/navbuttons_view"
diff --git a/quickstep/res/values/config.xml b/quickstep/res/values/config.xml
index e8cb5d5..41b2384 100644
--- a/quickstep/res/values/config.xml
+++ b/quickstep/res/values/config.xml
@@ -37,8 +37,8 @@
     <string name="taskbar_edu_tooltip_controller_class" translatable="false">com.android.launcher3.taskbar.TaskbarEduTooltipController</string>
     <string name="contextual_edu_manager_class" translatable="false">com.android.quickstep.contextualeducation.SystemContextualEduStatsManager</string>
     <string name="nav_handle_long_press_handler_class" translatable="false"></string>
-    <string name="assist_utils_class" translatable="false"></string>
-    <string name="assist_state_manager_class" translatable="false"></string>
+    <string name="contextual_search_invoker_class" translatable="false"></string>
+    <string name="contextual_search_state_manager_class" translatable="false"></string>
     <string name="api_wrapper_class" translatable="false">com.android.launcher3.uioverrides.SystemApiWrapper</string>
 
     <!-- The number of thumbnails and icons to keep in the cache. The thumbnail cache size also
diff --git a/quickstep/res/values/strings.xml b/quickstep/res/values/strings.xml
index 008766b..026e25c 100644
--- a/quickstep/res/values/strings.xml
+++ b/quickstep/res/values/strings.xml
@@ -28,6 +28,8 @@
     <string name="recent_task_option_freeform">Freeform</string>
     <!-- Title and content description for an option to enter desktop windowing mode for a given app -->
     <string name="recent_task_option_desktop">Desktop</string>
+    <!-- Title and content description for an option to move app to external display. -->
+    <string name="recent_task_option_external_display">Move to external display</string>
 
     <!-- Title and content description for Desktop tile in Recents screen that contains apps opened inside desktop windowing mode [CHAR LIMIT=NONE] -->
     <string name="recent_task_desktop">Desktop</string>
@@ -360,4 +362,8 @@
     <string name="bubble_bar_accessibility_announce_expand">expand <xliff:g id="bubble_description" example="some title from Messages">%1$s</xliff:g></string>
     <!-- Accessibility announcement when the bubble bar collapses. [CHAR LIMIT=NONE]-->
     <string name="bubble_bar_accessibility_announce_collapse">collapse <xliff:g id="bubble_description" example="some title from Messages">%1$s</xliff:g></string>
+
+    <!-- Name of Google's new feature to circle to search anything on your phone screen, without
+     switching apps. [CHAR_LIMIT=60] -->
+    <string name="search_gesture_feature_title">Circle to Search</string>
 </resources>
diff --git a/quickstep/src/com/android/launcher3/desktop/DesktopRecentsTransitionController.kt b/quickstep/src/com/android/launcher3/desktop/DesktopRecentsTransitionController.kt
index 8b3a032..ac1ffa6 100644
--- a/quickstep/src/com/android/launcher3/desktop/DesktopRecentsTransitionController.kt
+++ b/quickstep/src/com/android/launcher3/desktop/DesktopRecentsTransitionController.kt
@@ -66,6 +66,11 @@
         systemUiProxy.moveToDesktop(taskId, transitionSource)
     }
 
+    /** Move task to external display from recents view */
+    fun moveToExternalDisplay(taskId: Int) {
+        systemUiProxy.moveToExternalDisplay(taskId)
+    }
+
     private class RemoteDesktopLaunchTransitionRunner(
         private val desktopTaskView: DesktopTaskView,
         private val animated: Boolean,
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index 4014f06..29e1f4e 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -77,6 +77,7 @@
 import com.android.launcher3.util.PersistedItemArray;
 import com.android.quickstep.logging.SettingsChangeLogger;
 import com.android.quickstep.logging.StatsLogCompatManager;
+import com.android.quickstep.util.ContextualSearchStateManager;
 import com.android.systemui.shared.system.SysUiStatsLog;
 
 import java.util.ArrayList;
@@ -209,6 +210,8 @@
     @Override
     public void workspaceLoadComplete() {
         super.workspaceLoadComplete();
+        // Initialize ContextualSearchStateManager.
+        ContextualSearchStateManager.INSTANCE.get(mContext);
         recreatePredictors();
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index 09dbeb6..042bc9a 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -36,6 +36,7 @@
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimatedFloat;
+import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.logging.InstanceId;
 import com.android.launcher3.logging.InstanceIdSequence;
 import com.android.launcher3.model.data.ItemInfo;
@@ -68,14 +69,17 @@
     public static final int ALL_APPS_PAGE_PROGRESS_INDEX = 1;
     public static final int WIDGETS_PAGE_PROGRESS_INDEX = 2;
     public static final int SYSUI_SURFACE_PROGRESS_INDEX = 3;
+    public static final int LAUNCHER_PAUSE_PROGRESS_INDEX = 4;
 
-    public static final int DISPLAY_PROGRESS_COUNT = 4;
+    public static final int DISPLAY_PROGRESS_COUNT = 5;
 
     private final AnimatedFloat mTaskbarInAppDisplayProgress = new AnimatedFloat(
             this::onInAppDisplayProgressChanged);
     private final MultiPropertyFactory<AnimatedFloat> mTaskbarInAppDisplayProgressMultiProp =
             new MultiPropertyFactory<>(mTaskbarInAppDisplayProgress,
                     AnimatedFloat.VALUE, DISPLAY_PROGRESS_COUNT, Float::max);
+    private final AnimatedFloat mLauncherPauseProgress = new AnimatedFloat(
+            this::launcherPauseProgressUpdate);
 
     private final QuickstepLauncher mLauncher;
     private final HomeVisibilityState mHomeState;
@@ -191,6 +195,33 @@
     }
 
     /**
+     * Called when Launcher Activity is paused/resumed.
+     * <p>
+     * To avoid UI clash between taskbar & bottom sheet, shift nav buttons down on launcher
+     * pause/resume at home.
+     * @param paused if launcher is currently paused.
+     */
+    public void onLauncherPausedOrResumed(boolean paused) {
+        if (!FeatureFlags.enableHomeTransitionListener()) {
+            onLauncherVisibilityChanged(mLauncher.hasBeenResumed());
+            return;
+        }
+
+        // Animate navbar iff pause/resume from home, NOT to/from app (avoid overriding existing
+        // animations).
+        boolean launcherPauseOrResumeFromHome = mHomeState.isHomeVisible() && mControllers
+                .taskbarAutohideSuspendController.isSuspendedForTransientTaskbarInLauncher();
+        if (launcherPauseOrResumeFromHome) {
+            mLauncherPauseProgress.animateToValue(paused ? 1.0f : 0.0f).start();
+        }
+    }
+
+    private void launcherPauseProgressUpdate() {
+        onTaskbarInAppDisplayProgressUpdate(
+                mLauncherPauseProgress.value, LAUNCHER_PAUSE_PROGRESS_INDEX);
+    }
+
+    /**
      * Should be called from onResume() and onPause(), and animates the Taskbar accordingly.
      */
     @Override
@@ -364,18 +395,20 @@
         }
         if (mControllers.uiController.isIconAlignedWithHotseat()
                 && !mTaskbarLauncherStateController.isAnimatingToLauncher()) {
-            // Only animate the nav buttons while home and not animating home, otherwise let
+            // Only animate nav button position while home and not animating home, otherwise let
             // the TaskbarViewController handle it.
             mControllers.navbarButtonsViewController
-                    .getTaskbarNavButtonTranslationYForInAppDisplay()
+                    .getNavButtonTranslationYForInAppDisplay()
                     .updateValue(mLauncher.getDeviceProfile().getTaskbarOffsetY()
                             * mTaskbarInAppDisplayProgress.value);
-            mControllers.navbarButtonsViewController
-                    .getOnTaskbarBackgroundNavButtonColorOverride().updateValue(progress);
+            if (!mLauncher.isPaused()) {
+                mControllers.navbarButtonsViewController
+                        .getOnTaskbarBackgroundNavButtonColorOverride().updateValue(progress);
+            }
         }
     }
 
-    /** Returns true iff any in-app display progress > 0. */
+    @Override
     public boolean shouldUseInAppLayout() {
         return mTaskbarInAppDisplayProgress.value > 0;
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 7d8e93c..cfcbd2f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -183,7 +183,7 @@
 
     private final AnimatedFloat mTaskbarNavButtonTranslationY = new AnimatedFloat(
             this::updateNavButtonTranslationY);
-    private final AnimatedFloat mTaskbarNavButtonTranslationYForInAppDisplay = new AnimatedFloat(
+    private final AnimatedFloat mNavButtonTranslationYForInAppDisplay = new AnimatedFloat(
             this::updateNavButtonTranslationY);
     private final AnimatedFloat mTaskbarNavButtonTranslationYForIme = new AnimatedFloat(
             this::updateNavButtonTranslationY);
@@ -704,8 +704,8 @@
     }
 
     /** Use to set the translationY for the all nav+contextual buttons when in Launcher */
-    public AnimatedFloat getTaskbarNavButtonTranslationYForInAppDisplay() {
-        return mTaskbarNavButtonTranslationYForInAppDisplay;
+    public AnimatedFloat getNavButtonTranslationYForInAppDisplay() {
+        return mNavButtonTranslationYForInAppDisplay;
     }
 
     /** Use to set the dark intensity for the all nav+contextual buttons */
@@ -751,21 +751,23 @@
         if (mContext.isPhoneButtonNavMode()) {
             return;
         }
-        final float normalTranslationY = mTaskbarNavButtonTranslationY.value;
-        final float imeAdjustmentTranslationY = mTaskbarNavButtonTranslationYForIme.value;
-        TaskbarUIController uiController = mControllers.uiController;
-        final float inAppDisplayAdjustmentTranslationY =
-                (uiController instanceof LauncherTaskbarUIController
-                        && ((LauncherTaskbarUIController) uiController).shouldUseInAppLayout())
-                        ? mTaskbarNavButtonTranslationYForInAppDisplay.value : 0;
-
-        mLastSetNavButtonTranslationY = normalTranslationY
-                + imeAdjustmentTranslationY
-                + inAppDisplayAdjustmentTranslationY;
+        mLastSetNavButtonTranslationY = calculateNavButtonTranslationY();
         mNavButtonsView.setTranslationY(mLastSetNavButtonTranslationY);
     }
 
     /**
+     * Calculates the translationY of the nav buttons based on the current device state.
+     */
+    private float calculateNavButtonTranslationY() {
+        float translationY =
+                mTaskbarNavButtonTranslationY.value + mTaskbarNavButtonTranslationYForIme.value;
+        if (mControllers.uiController.shouldUseInAppLayout()) {
+            translationY += mNavButtonTranslationYForInAppDisplay.value;
+        }
+        return translationY;
+    }
+
+    /**
      * Sets Taskbar 3-button mode icon colors based on the
      * {@link #mTaskbarNavButtonDarkIntensity} value piped in from Framework. For certain cases
      * in large screen taskbar where there may be opaque surfaces, the selected SystemUI button
@@ -1162,7 +1164,7 @@
         pw.println(prefix + "\t\tmTaskbarNavButtonTranslationY="
                 + mTaskbarNavButtonTranslationY.value);
         pw.println(prefix + "\t\tmTaskbarNavButtonTranslationYForInAppDisplay="
-                + mTaskbarNavButtonTranslationYForInAppDisplay.value);
+                + mNavButtonTranslationYForInAppDisplay.value);
         pw.println(prefix + "\t\tmTaskbarNavButtonTranslationYForIme="
                 + mTaskbarNavButtonTranslationYForIme.value);
         pw.println(prefix + "\t\tmTaskbarNavButtonDarkIntensity="
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 24e4d2b..1b9614a 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -71,6 +71,7 @@
 import android.view.View;
 import android.view.WindowInsets;
 import android.view.WindowManager;
+import android.widget.FrameLayout;
 import android.widget.Toast;
 import android.window.RemoteTransition;
 
@@ -268,8 +269,10 @@
         NearestTouchFrame navButtonsView = mDragLayer.findViewById(R.id.navbuttons_view);
         StashedHandleView stashedHandleView = mDragLayer.findViewById(R.id.stashed_handle);
         BubbleBarView bubbleBarView = null;
+        FrameLayout bubbleBarContainer = null;
         if (isTransientTaskbar || Flags.enableBubbleBarInPersistentTaskBar()) {
             bubbleBarView = mDragLayer.findViewById(R.id.taskbar_bubbles);
+            bubbleBarContainer = mDragLayer.findViewById(R.id.taskbar_bubbles_container);
         }
         StashedHandleView bubbleHandleView = mDragLayer.findViewById(R.id.stashed_bubble_handle);
 
@@ -296,7 +299,7 @@
                     : new PersistentBubbleStashController(dimensionsProvider);
             bubbleControllersOptional = Optional.of(new BubbleControllers(
                     new BubbleBarController(this, bubbleBarView),
-                    new BubbleBarViewController(this, bubbleBarView),
+                    new BubbleBarViewController(this, bubbleBarView, bubbleBarContainer),
                     bubbleStashController,
                     bubbleHandleController,
                     new BubbleDragController(this),
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
index 219a24a..4a85acc 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -195,11 +195,12 @@
         };
 
         if (taskbarDesktopModeController.getAreDesktopTasksVisible()) {
-            mCornerRoundness.updateValue(taskbarDesktopModeController.getTaskbarCornerRoundness(
-                    mSharedState.showCornerRadiusInDesktopMode));
+            mCornerRoundness.value = taskbarDesktopModeController.getTaskbarCornerRoundness(
+                    mSharedState.showCornerRadiusInDesktopMode);
         } else {
-            mCornerRoundness.updateValue(TaskbarBackgroundRenderer.MAX_ROUNDNESS);
+            mCornerRoundness.value = TaskbarBackgroundRenderer.MAX_ROUNDNESS;
         }
+        updateCornerRoundness();
         onPostInit();
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduTooltipController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduTooltipController.kt
index 06376d3..ade8f8c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduTooltipController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduTooltipController.kt
@@ -49,6 +49,7 @@
 import com.android.launcher3.util.ResourceBasedOverride
 import com.android.launcher3.views.ActivityContext
 import com.android.launcher3.views.BaseDragLayer
+import com.android.quickstep.util.ContextualSearchInvoker
 import com.android.quickstep.util.LottieAnimationColorUtils
 import java.io.PrintWriter
 
@@ -80,7 +81,11 @@
     ResourceBasedOverride, LoggableTaskbarController {
 
     protected val activityContext: TaskbarActivityContext = ActivityContext.lookupContext(context)
-    open val shouldShowSearchEdu = false
+    open val shouldShowSearchEdu: Boolean
+        get() =
+            ContextualSearchInvoker.newInstance(activityContext)
+                .runContextualSearchInvocationChecksAndLogFailures()
+
     private val isTooltipEnabled: Boolean
         get() {
             return !Utilities.isRunningInTestHarness() &&
@@ -351,19 +356,19 @@
             overlayContext.layoutInflater.inflate(
                 R.layout.taskbar_edu_tooltip,
                 overlayContext.dragLayer,
-                false
+                false,
             ) as TaskbarEduTooltip
 
         controllers.taskbarAutohideSuspendController.updateFlag(
             FLAG_AUTOHIDE_SUSPEND_EDU_OPEN,
-            true
+            true,
         )
 
         tooltip.onCloseCallback = {
             this.tooltip = null
             controllers.taskbarAutohideSuspendController.updateFlag(
                 FLAG_AUTOHIDE_SUSPEND_EDU_OPEN,
-                false
+                false,
             )
             controllers.taskbarStashController.updateAndAnimateTransientTaskbar(true)
         }
@@ -378,7 +383,7 @@
             override fun performAccessibilityAction(
                 host: View,
                 action: Int,
-                args: Bundle?
+                args: Bundle?,
             ): Boolean {
                 if (action == R.id.close) {
                     hide()
@@ -396,13 +401,13 @@
 
             override fun onInitializeAccessibilityNodeInfo(
                 host: View,
-                info: AccessibilityNodeInfo
+                info: AccessibilityNodeInfo,
             ) {
                 super.onInitializeAccessibilityNodeInfo(host, info)
                 info.addAction(
                     AccessibilityNodeInfo.AccessibilityAction(
                         R.id.close,
-                        host.context?.getText(R.string.taskbar_edu_close)
+                        host.context?.getText(R.string.taskbar_edu_close),
                     )
                 )
             }
@@ -421,7 +426,7 @@
             return ResourceBasedOverride.Overrides.getObject(
                 TaskbarEduTooltipController::class.java,
                 context,
-                R.string.taskbar_edu_tooltip_controller_class
+                R.string.taskbar_edu_tooltip_controller_class,
             )
         }
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index 78e7b47..c18cf28 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -16,7 +16,6 @@
 package com.android.launcher3.taskbar;
 
 import static android.content.Context.RECEIVER_NOT_EXPORTED;
-import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR;
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR_PANEL;
 
@@ -72,7 +71,7 @@
 import com.android.quickstep.AllAppsActionManager;
 import com.android.quickstep.RecentsActivity;
 import com.android.quickstep.SystemUiProxy;
-import com.android.quickstep.util.AssistUtils;
+import com.android.quickstep.util.ContextualSearchInvoker;
 import com.android.systemui.shared.statusbar.phone.BarTransitions;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.shared.system.QuickStepContract.SystemUiStateFlags;
@@ -220,7 +219,7 @@
             TaskbarNavButtonCallbacks navCallbacks,
             @NonNull DesktopVisibilityController desktopVisibilityController) {
         Display display =
-                context.getSystemService(DisplayManager.class).getDisplay(DEFAULT_DISPLAY);
+                context.getSystemService(DisplayManager.class).getDisplay(context.getDisplayId());
         mContext = context.createWindowContext(display,
                 ENABLE_TASKBAR_NAVBAR_UNIFICATION ? TYPE_NAVIGATION_BAR : TYPE_NAVIGATION_BAR_PANEL,
                 null);
@@ -250,7 +249,7 @@
                 SystemUiProxy.INSTANCE.get(mContext),
                 ContextualEduStatsManager.INSTANCE.get(mContext),
                 new Handler(),
-                AssistUtils.newInstance(mContext));
+                ContextualSearchInvoker.newInstance(mContext));
         mComponentCallbacks = new ComponentCallbacks() {
             private Configuration mOldConfig = mContext.getResources().getConfiguration();
 
@@ -672,11 +671,6 @@
     @VisibleForTesting
     public void setSuspended(boolean isSuspended) {
         mIsSuspended = isSuspended;
-        if (mIsSuspended) {
-            removeTaskbarRootViewFromWindow();
-        } else {
-            addTaskbarRootViewToWindow();
-        }
     }
 
     private void addTaskbarRootViewToWindow() {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
index 15c35b6..8947914 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
@@ -51,7 +51,7 @@
 import com.android.launcher3.testing.shared.TestProtocol;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TaskUtils;
-import com.android.quickstep.util.AssistUtils;
+import com.android.quickstep.util.ContextualSearchInvoker;
 import com.android.systemui.contextualeducation.GestureType;
 import com.android.systemui.shared.system.QuickStepContract.SystemUiStateFlags;
 
@@ -113,7 +113,7 @@
     private final SystemUiProxy mSystemUiProxy;
     private final ContextualEduStatsManager mContextualEduStatsManager;
     private final Handler mHandler;
-    private final AssistUtils mAssistUtils;
+    private final ContextualSearchInvoker mContextualSearchInvoker;
     @Nullable private StatsLogManager mStatsLogManager;
 
     private final Runnable mResetLongPress = this::resetScreenUnpin;
@@ -124,13 +124,13 @@
             SystemUiProxy systemUiProxy,
             ContextualEduStatsManager contextualEduStatsManager,
             Handler handler,
-            AssistUtils assistUtils) {
+            ContextualSearchInvoker contextualSearchInvoker) {
         mContext = context;
         mCallbacks = callbacks;
         mSystemUiProxy = systemUiProxy;
         mContextualEduStatsManager = contextualEduStatsManager;
         mHandler = handler;
-        mAssistUtils = assistUtils;
+        mContextualSearchInvoker = contextualSearchInvoker;
     }
 
     public void onButtonClick(@TaskbarButton int buttonType, View view) {
@@ -344,8 +344,9 @@
         if (mScreenPinned || !mAssistantLongPressEnabled) {
             return;
         }
-        // Attempt to start Assist with AssistUtils, otherwise fall back to SysUi's implementation.
-        if (!mAssistUtils.tryStartAssistOverride(INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS)) {
+        // Attempt to start Contextual Search, otherwise fall back to SysUi's implementation.
+        if (!mContextualSearchInvoker.tryStartAssistOverride(
+                INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS)) {
             Bundle args = new Bundle();
             args.putInt(INVOCATION_TYPE_KEY, INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS);
             mSystemUiProxy.startAssistant(args);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
index b80aaf8..7030088 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
@@ -91,6 +91,14 @@
     protected void onStashedInAppChanged() { }
 
     /**
+     * Whether the Taskbar should use in-app layout.
+     * @return {@code true} iff in-app display progress > 0 or Launcher Activity paused.
+     */
+    public boolean shouldUseInAppLayout() {
+        return false;
+    }
+
+    /**
      * Called when taskbar icon layout bounds change.
      */
     protected void onIconLayoutBoundsChanged() { }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index 8763509..70dcfe8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -192,13 +192,44 @@
     }
 
     /**
-     // @return the maximum number of 'icons' that can fit in the taskbar.
-     // TODO(368119679): Assumes that they are all the same size.
+     * @return the maximum number of 'icons' that can fit in the taskbar.
      */
     private int calculateMaxNumIcons() {
-        int availableWidth = mActivityContext.getDeviceProfile().widthPx
-                - (mActivityContext.getDeviceProfile().edgeMarginPx * 2);
-        return Math.floorDiv(availableWidth, mIconTouchSize);
+        DeviceProfile deviceProfile = mActivityContext.getDeviceProfile();
+        int availableWidth = deviceProfile.widthPx;
+
+        // Reserve space required for edge margins, or for navbar if shown. If task bar needs to be
+        // center aligned with nav bar shown, reserve space on both sides.
+        availableWidth -= Math.max(deviceProfile.edgeMarginPx, deviceProfile.hotseatBarEndOffset);
+        availableWidth -= Math.max(deviceProfile.edgeMarginPx,
+                mShouldTryStartAlign ? 0 : deviceProfile.hotseatBarEndOffset);
+
+        // The space taken by an item icon used during layout.
+        int iconSize = 2 * mItemMarginLeftRight + mIconTouchSize;
+
+        int additionalIcons = 0;
+
+        if (mTaskbarDividerContainer != null) {
+            // Space for divider icon is reduced during layout compared to normal icon size, reserve
+            // space for the divider separately.
+            availableWidth -= iconSize - 4 * mItemMarginLeftRight;
+            ++additionalIcons;
+        }
+
+        // All apps icon takes less space compared to normal icon size, reserve space for the icon
+        // separately.
+        if (mAllAppsButtonContainer != null) {
+            boolean forceTransientTaskbarSize =
+                    enableTaskbarPinning() && !mActivityContext.isThreeButtonNav();
+            availableWidth -= iconSize - (int) getResources().getDimension(
+                    mAllAppsButtonContainer.getAllAppsButtonTranslationXOffset(
+                            forceTransientTaskbarSize || (
+                                    DisplayController.isTransientTaskbar(mActivityContext)
+                                            && !mActivityContext.isPhoneMode())));
+            ++additionalIcons;
+        }
+
+        return Math.floorDiv(availableWidth, iconSize) + additionalIcons;
     }
 
     @Override
@@ -412,13 +443,33 @@
         if (mTaskbarDividerContainer != null && !recentTasks.isEmpty()) {
             addView(mTaskbarDividerContainer, nextViewIndex++);
             mAddedDividerForRecents = true;
-            if (mTaskbarOverflowView != null) {
+        }
+
+        // At this point, the all apps button has not been added as a child view, but needs to be
+        // accounted for when comparing current icon count to max number of icons.
+        int nonTaskIconsToBeAdded = 1;
+
+        boolean supportsOverflow = Flags.taskbarOverflow();
+        if (supportsOverflow) {
+            int numberOfSupportedRecents = 0;
+            for (GroupTask task : recentTasks) {
+                // TODO(b/343289567 and b/316004172): support app pairs and desktop mode.
+                if (!task.hasMultipleTasks()) {
+                    ++numberOfSupportedRecents;
+                }
+            }
+            if (nextViewIndex + numberOfSupportedRecents + nonTaskIconsToBeAdded > mMaxNumIcons
+                    && mTaskbarOverflowView != null) {
                 addView(mTaskbarOverflowView, nextViewIndex++);
             }
         }
 
         // Add Recent/Running icons.
         for (GroupTask task : recentTasks) {
+            if (supportsOverflow && nextViewIndex + nonTaskIconsToBeAdded >= mMaxNumIcons) {
+                break;
+            }
+
             // Replace any Recent views with the appropriate type if it's not already that type.
             final int expectedLayoutResId;
             boolean isCollection = false;
@@ -488,8 +539,6 @@
             }
         }
 
-        updateRecentAppsToFit();
-
         if (mActivityContext.getDeviceProfile().isQsbInline) {
             addView(mQsb, mIsRtl ? getChildCount() : 0);
             // Always set QSB to invisible after re-adding.
@@ -497,45 +546,6 @@
         }
     }
 
-    /**
-     * Updates the recent apps portion of the taskbar by:
-     * - Removing overflow affordance if overflow is not needed.
-     * - Removing any recent apps that do not fit.
-     */
-    private void updateRecentAppsToFit() {
-        if (!Flags.taskbarOverflow()) {
-            return;
-        }
-        int indexOfFirstRecentApp = -1;
-        int size = getChildCount();
-        boolean removeOverflowView = true;
-
-        for (int i = 0; i < size; ++i) {
-            if (getChildAt(i).getTag() instanceof GroupTask) {
-                indexOfFirstRecentApp = i;
-                removeOverflowView = false;
-                break;
-            }
-        }
-
-        if (indexOfFirstRecentApp != -1) {
-            // We pre-maturely added the overflow icon, so we can take it out of the count.
-            int numRecentAppsToRemove = Math.max(0, getChildCount() - mMaxNumIcons + 1);
-            if (numRecentAppsToRemove <= 1) {
-                // We can fit all of the recent apps if we remove the overflow icon.
-                removeOverflowView = true;
-            } else {
-                for (int i = 0; i < numRecentAppsToRemove; ++i) {
-                    removeAndRecycle(getChildAt(indexOfFirstRecentApp));
-                }
-            }
-        }
-
-        if (removeOverflowView) {
-            removeView(mTaskbarOverflowView);
-        }
-    }
-
     /** Binds the GroupTask to the BubbleTextView to be ready to present to the user. */
     public void applyGroupTaskToBubbleTextView(BubbleTextView btv, GroupTask groupTask) {
         // TODO(b/343289567): support app pairs.
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacks.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacks.java
index 176be1c..8bc1e12 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacks.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacks.java
@@ -19,6 +19,7 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_ALLAPPS_BUTTON_LONG_PRESS;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_ALLAPPS_BUTTON_TAP;
 
+import android.content.Context;
 import android.view.InputDevice;
 import android.view.MotionEvent;
 import android.view.View;
@@ -64,7 +65,8 @@
         mActivity.getStatsLogManager().logger().log(LAUNCHER_TASKBAR_ALLAPPS_BUTTON_LONG_PRESS);
     }
 
-    public boolean isAllAppsButtonHapticFeedbackEnabled() {
+    /** @return true if haptic feedback should occur when long pressing the all apps button. */
+    public boolean isAllAppsButtonHapticFeedbackEnabled(Context context) {
         return false;
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacksFactory.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacksFactory.kt
index ba0f5a0..704d6cf 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacksFactory.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacksFactory.kt
@@ -16,10 +16,14 @@
 
 package com.android.launcher3.taskbar
 
+import android.app.contextualsearch.ContextualSearchManager.ENTRYPOINT_LONG_PRESS_META
 import android.content.Context
 import com.android.launcher3.R
+import com.android.launcher3.logging.StatsLogManager
 import com.android.launcher3.util.ResourceBasedOverride
 import com.android.launcher3.util.ResourceBasedOverride.Overrides
+import com.android.quickstep.TopTaskTracker
+import com.android.quickstep.util.ContextualSearchInvoker
 
 /** Creates [TaskbarViewCallbacks] instances. */
 open class TaskbarViewCallbacksFactory : ResourceBasedOverride {
@@ -28,7 +32,35 @@
         activity: TaskbarActivityContext,
         controllers: TaskbarControllers,
         taskbarView: TaskbarView,
-    ): TaskbarViewCallbacks = TaskbarViewCallbacks(activity, controllers, taskbarView)
+    ): TaskbarViewCallbacks {
+        return object : TaskbarViewCallbacks(activity, controllers, taskbarView) {
+            override fun triggerAllAppsButtonLongClick() {
+                super.triggerAllAppsButtonLongClick()
+
+                val contextualSearchInvoked =
+                    ContextualSearchInvoker.newInstance(activity).show(ENTRYPOINT_LONG_PRESS_META)
+                if (contextualSearchInvoked) {
+                    val runningPackage =
+                        TopTaskTracker.INSTANCE[activity].getCachedTopTask(
+                                /* filterOnlyVisibleRecents */ true
+                            )
+                            .getPackageName()
+                    activity.statsLogManager
+                        .logger()
+                        .withPackageName(runningPackage)
+                        .log(StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_SUCCESSFUL_META)
+                }
+            }
+
+            override fun isAllAppsButtonHapticFeedbackEnabled(context: Context): Boolean {
+                return longPressAllAppsToStartContextualSearch(context)
+            }
+        }
+    }
+
+    open fun longPressAllAppsToStartContextualSearch(context: Context): Boolean =
+        ContextualSearchInvoker.newInstance(context)
+            .runContextualSearchInvocationChecksAndLogFailures()
 
     companion object {
         @JvmStatic
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
index bc61c72..c4d9e50 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
@@ -234,7 +234,7 @@
         mTaskbarNavButtonTranslationY =
                 controllers.navbarButtonsViewController.getTaskbarNavButtonTranslationY();
         mTaskbarNavButtonTranslationYForInAppDisplay = controllers.navbarButtonsViewController
-                .getTaskbarNavButtonTranslationYForInAppDisplay();
+                .getNavButtonTranslationYForInAppDisplay();
 
         mActivity.addOnDeviceProfileChangeListener(mDeviceProfileChangeListener);
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index da9dcf7..63f101f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -29,6 +29,7 @@
 import android.util.TypedValue;
 import android.view.MotionEvent;
 import android.view.View;
+import android.widget.FrameLayout;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -42,6 +43,7 @@
 import com.android.launcher3.taskbar.TaskbarInsetsController;
 import com.android.launcher3.taskbar.TaskbarStashController;
 import com.android.launcher3.taskbar.bubbles.animation.BubbleBarViewAnimator;
+import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutController;
 import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutPositioner;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController;
 import com.android.launcher3.util.MultiPropertyFactory;
@@ -116,15 +118,19 @@
     public boolean mOverflowAdded;
 
     private BubbleBarViewAnimator mBubbleBarViewAnimator;
+    private final FrameLayout mBubbleBarContainer;
+    private BubbleBarFlyoutController mBubbleBarFlyoutController;
 
     private final TimeSource mTimeSource = System::currentTimeMillis;
 
     @Nullable
     private BubbleBarBoundsChangeListener mBoundsChangeListener;
 
-    public BubbleBarViewController(TaskbarActivityContext activity, BubbleBarView barView) {
+    public BubbleBarViewController(TaskbarActivityContext activity, BubbleBarView barView,
+            FrameLayout bubbleBarContainer) {
         mActivity = activity;
         mBarView = barView;
+        mBubbleBarContainer = bubbleBarContainer;
         mSystemUiProxy = SystemUiProxy.INSTANCE.get(mActivity);
         mBubbleBarAlpha = new MultiValueAlpha(mBarView, 1 /* num alpha channels */);
         mIconSize = activity.getResources().getDimensionPixelSize(
@@ -139,6 +145,8 @@
         mBubbleDragController = bubbleControllers.bubbleDragController;
         mTaskbarStashController = controllers.taskbarStashController;
         mTaskbarInsetsController = controllers.taskbarInsetsController;
+        mBubbleBarFlyoutController = new BubbleBarFlyoutController(
+                mBubbleBarContainer, createFlyoutPositioner(), createFlyoutTopBoundaryListener());
         mBubbleBarViewAnimator = new BubbleBarViewAnimator(
                 mBarView, mBubbleStashController, mBubbleBarController::showExpandedView);
         mTaskbarViewPropertiesProvider = taskbarViewPropertiesProvider;
@@ -264,6 +272,21 @@
         };
     }
 
+    private BubbleBarFlyoutController.TopBoundaryListener createFlyoutTopBoundaryListener() {
+        return new BubbleBarFlyoutController.TopBoundaryListener() {
+            @Override
+            public void extendTopBoundary(int space) {
+                int defaultSize = mActivity.getDefaultTaskbarWindowSize();
+                mActivity.setTaskbarWindowSize(defaultSize + space);
+            }
+
+            @Override
+            public void resetTopBoundary() {
+                mActivity.setTaskbarWindowSize(mActivity.getDefaultTaskbarWindowSize());
+            }
+        };
+    }
+
     private void onBubbleClicked(BubbleView bubbleView) {
         bubbleView.markSeen();
         BubbleBarItem bubble = bubbleView.getBubble();
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
index 49760ff..c431deb 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
@@ -21,20 +21,30 @@
 import android.widget.FrameLayout
 import androidx.core.animation.ValueAnimator
 import com.android.launcher3.R
+import com.android.systemui.util.doOnEnd
+import com.android.systemui.util.doOnStart
 
 /** Creates and manages the visibility of the [BubbleBarFlyoutView]. */
-class BubbleBarFlyoutController(
+class BubbleBarFlyoutController
+@JvmOverloads
+constructor(
     private val container: FrameLayout,
     private val positioner: BubbleBarFlyoutPositioner,
+    private val topBoundaryListener: TopBoundaryListener,
+    private val flyoutScheduler: FlyoutScheduler = HandlerScheduler(container),
 ) {
 
+    private companion object {
+        const val EXPAND_COLLAPSE_ANIMATION_DURATION_MS = 250L
+    }
+
     private var flyout: BubbleBarFlyoutView? = null
     private val horizontalMargin =
         container.context.resources.getDimensionPixelSize(R.dimen.transient_taskbar_bottom_margin)
 
     fun setUpFlyout(message: BubbleBarFlyoutMessage) {
         flyout?.let(container::removeView)
-        val flyout = BubbleBarFlyoutView(container.context, positioner)
+        val flyout = BubbleBarFlyoutView(container.context, positioner, flyoutScheduler)
 
         flyout.translationY = positioner.targetTy
 
@@ -48,17 +58,43 @@
         lp.marginEnd = horizontalMargin
         container.addView(flyout, lp)
 
-        val animator = ValueAnimator.ofFloat(0f, 1f)
+        val animator =
+            ValueAnimator.ofFloat(0f, 1f).setDuration(EXPAND_COLLAPSE_ANIMATION_DURATION_MS)
         animator.addUpdateListener { _ ->
             flyout.updateExpansionProgress(animator.animatedValue as Float)
         }
+        animator.doOnStart {
+            val flyoutTop = flyout.top + flyout.translationY
+            // If the top position of the flyout is negative, then it's bleeding over the
+            // top boundary of its parent view
+            if (flyoutTop < 0) topBoundaryListener.extendTopBoundary(space = -flyoutTop.toInt())
+        }
         flyout.showFromCollapsed(message) { animator.start() }
         this.flyout = flyout
     }
 
-    fun hideFlyout() {
+    fun hideFlyout(endAction: () -> Unit) {
         val flyout = this.flyout ?: return
-        container.removeView(flyout)
-        this.flyout = null
+        val animator =
+            ValueAnimator.ofFloat(1f, 0f).setDuration(EXPAND_COLLAPSE_ANIMATION_DURATION_MS)
+        animator.addUpdateListener { _ ->
+            flyout.updateExpansionProgress(animator.animatedValue as Float)
+        }
+        animator.doOnEnd {
+            container.removeView(flyout)
+            this@BubbleBarFlyoutController.flyout = null
+            topBoundaryListener.resetTopBoundary()
+            endAction()
+        }
+        animator.start()
+    }
+
+    /** Notifies when the top boundary of the flyout view changes. */
+    interface TopBoundaryListener {
+        /** Requests to extend the top boundary of the parent to fully include the flyout. */
+        fun extendTopBoundary(space: Int)
+
+        /** Resets the top boundary of the parent. */
+        fun resetTopBoundary()
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
index 2022a42..c60fba2 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
@@ -36,14 +36,18 @@
 import com.android.launcher3.popup.RoundedArrowDrawable
 
 /** The flyout view used to notify the user of a new bubble notification. */
-class BubbleBarFlyoutView(context: Context, private val positioner: BubbleBarFlyoutPositioner) :
-    ConstraintLayout(context) {
+class BubbleBarFlyoutView(
+    context: Context,
+    private val positioner: BubbleBarFlyoutPositioner,
+    scheduler: FlyoutScheduler? = null,
+) : ConstraintLayout(context) {
 
     private companion object {
         // the minimum progress of the expansion animation before the content starts fading in.
         const val MIN_EXPANSION_PROGRESS_FOR_CONTENT_ALPHA = 0.75f
     }
 
+    private val scheduler: FlyoutScheduler = scheduler ?: HandlerScheduler(this)
     private val title: TextView by
         lazy(LazyThreadSafetyMode.NONE) { findViewById(R.id.bubble_flyout_title) }
 
@@ -197,11 +201,10 @@
 
         // post the request to start the expand animation to the looper so the view can measure
         // itself
-        post(expandAnimation)
+        scheduler.runAfterLayout(expandAnimation)
     }
 
     private fun setData(flyoutMessage: BubbleBarFlyoutMessage) {
-        // the avatar is only displayed in group chat messages
         if (flyoutMessage.icon != null) {
             icon.visibility = VISIBLE
             icon.setImageDrawable(flyoutMessage.icon)
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/FlyoutScheduler.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/FlyoutScheduler.kt
new file mode 100644
index 0000000..6f5d700
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/FlyoutScheduler.kt
@@ -0,0 +1,32 @@
+/*
+ * 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.bubbles.flyout
+
+import android.view.View
+
+/** Interface for scheduling jobs by flyout. */
+fun interface FlyoutScheduler {
+    /** Runs the given [block] after layout. */
+    fun runAfterLayout(block: () -> Unit)
+}
+
+/** A [FlyoutScheduler] that uses a Handler to schedule jobs. */
+class HandlerScheduler(val view: View) : FlyoutScheduler {
+    override fun runAfterLayout(block: () -> Unit) {
+        view.post(block)
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt
index e6c0b2f..c5f8aa0 100644
--- a/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt
@@ -37,7 +37,7 @@
 import com.android.launcher3.views.ActivityContext
 import com.android.launcher3.views.IconButtonView
 import com.android.quickstep.DeviceConfigWrapper
-import com.android.quickstep.util.AssistStateManager
+import com.android.quickstep.util.ContextualSearchStateManager
 
 /** Taskbar all apps button container for customizable taskbar. */
 class TaskbarAllAppsButtonContainer
@@ -79,17 +79,18 @@
         setOnClickListener(this::onAllAppsButtonClick)
         setOnLongClickListener(this::onAllAppsButtonLongClick)
         setOnTouchListener(this::onAllAppsButtonTouch)
-        isHapticFeedbackEnabled = taskbarViewCallbacks.isAllAppsButtonHapticFeedbackEnabled()
+        isHapticFeedbackEnabled =
+            taskbarViewCallbacks.isAllAppsButtonHapticFeedbackEnabled(mContext)
         allAppsTouchRunnable = Runnable {
             taskbarViewCallbacks.triggerAllAppsButtonLongClick()
             allAppsTouchTriggered = true
         }
-        val assistStateManager = AssistStateManager.INSTANCE[mContext]
+        val contextualSearchStateManager = ContextualSearchStateManager.INSTANCE[mContext]
         if (
             DeviceConfigWrapper.get().customLpaaThresholds &&
-                assistStateManager.lpnhDurationMillis.isPresent
+                contextualSearchStateManager.lpnhDurationMillis.isPresent
         ) {
-            allAppsButtonTouchDelayMs = assistStateManager.lpnhDurationMillis.get()
+            allAppsButtonTouchDelayMs = contextualSearchStateManager.lpnhDurationMillis.get()
         }
     }
 
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 39bf6ac..ce4e980 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -419,10 +419,8 @@
             mDepthController.setActivityStarted(isStarted());
         }
 
-        if ((changeBits & ACTIVITY_STATE_RESUMED) != 0) {
-            if (!FeatureFlags.enableHomeTransitionListener() && mTaskbarUIController != null) {
-                mTaskbarUIController.onLauncherVisibilityChanged(hasBeenResumed());
-            }
+        if ((changeBits & ACTIVITY_STATE_RESUMED) != 0 && mTaskbarUIController != null) {
+            mTaskbarUIController.onLauncherPausedOrResumed(isPaused());
         }
 
         super.onActivityFlagsChanged(changeBits);
diff --git a/quickstep/src/com/android/quickstep/BaseContainerInterface.java b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
index 7786353..143ef12 100644
--- a/quickstep/src/com/android/quickstep/BaseContainerInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
@@ -378,9 +378,6 @@
     public static void getTaskDimension(Context context, DeviceProfile dp, PointF out) {
         out.x = dp.widthPx;
         out.y = dp.heightPx;
-        if (dp.isTablet && !DisplayController.isTransientTaskbar(context)) {
-            out.y -= dp.taskbarHeight;
-        }
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt b/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt
new file mode 100644
index 0000000..46c4f36
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt
@@ -0,0 +1,86 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep
+
+import android.view.View
+import com.android.launcher3.AbstractFloatingViewHelper
+import com.android.launcher3.R
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent
+import com.android.launcher3.popup.SystemShortcut
+import com.android.quickstep.views.RecentsView
+import com.android.quickstep.views.RecentsViewContainer
+import com.android.quickstep.views.TaskContainer
+import com.android.window.flags.Flags
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
+
+/** A menu item that allows the user to move the current app into external display. */
+class ExternalDisplaySystemShortcut(
+    container: RecentsViewContainer,
+    abstractFloatingViewHelper: AbstractFloatingViewHelper,
+    private val taskContainer: TaskContainer,
+) :
+    SystemShortcut<RecentsViewContainer>(
+        R.drawable.ic_external_display,
+        R.string.recent_task_option_external_display,
+        container,
+        taskContainer.itemInfo,
+        taskContainer.taskView,
+        abstractFloatingViewHelper,
+    ) {
+    override fun onClick(view: View) {
+        dismissTaskMenuView()
+        val recentsView = mTarget.getOverviewPanel<RecentsView<*, *>>()
+        recentsView.moveTaskToExternalDisplay(taskContainer) {
+            mTarget.statsLogManager
+                .logger()
+                .withItemInfo(taskContainer.itemInfo)
+                .log(LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_EXTERNAL_DISPLAY_TAP)
+        }
+    }
+
+    companion object {
+        @JvmOverloads
+        /**
+         * Creates a factory for creating move task to external display system shortcuts in
+         * [com.android.quickstep.TaskOverlayFactory].
+         */
+        fun createFactory(
+            abstractFloatingViewHelper: AbstractFloatingViewHelper = AbstractFloatingViewHelper()
+        ): TaskShortcutFactory =
+            object : TaskShortcutFactory {
+                override fun getShortcuts(
+                    container: RecentsViewContainer,
+                    taskContainer: TaskContainer,
+                ): List<ExternalDisplaySystemShortcut>? {
+                    return if (
+                        DesktopModeStatus.canEnterDesktopMode(container.asContext()) &&
+                            Flags.moveToExternalDisplayShortcut()
+                    )
+                        listOf(
+                            ExternalDisplaySystemShortcut(
+                                container,
+                                abstractFloatingViewHelper,
+                                taskContainer,
+                            )
+                        )
+                    else null
+                }
+
+                override fun showForGroupedTask() = true
+            }
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
index 461f963..e23947b 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
@@ -53,6 +53,7 @@
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper
 import java.io.PrintWriter
 import java.util.concurrent.ConcurrentLinkedDeque
+import java.util.concurrent.Executor
 import kotlin.coroutines.resume
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.SupervisorJob
@@ -69,6 +70,7 @@
     private val overviewComponentObserver: OverviewComponentObserver,
     private val taskAnimationManager: TaskAnimationManager,
     private val dispatcherProvider: DispatcherProvider = ProductionDispatchers,
+    private val uiExecutor: Executor = Executors.MAIN_EXECUTOR,
 ) {
     private val coroutineScope = CoroutineScope(SupervisorJob() + dispatcherProvider.default)
 
@@ -85,7 +87,7 @@
         get() = overviewComponentObserver.containerInterface
 
     private val visibleRecentsView: RecentsView<*, *>?
-        get() = containerInterface.getVisibleRecentsView<RecentsView<*, *>>()
+        get() = containerInterface.getVisibleRecentsView()
 
     /**
      * Adds a command to be executed next, after all pending tasks are completed. Max commands that
@@ -105,11 +107,7 @@
 
         if (commandQueue.size == 1) {
             Log.d(TAG, "execute: $command - queue size: ${commandQueue.size}")
-            if (enableOverviewCommandHelperTimeout()) {
-                coroutineScope.launch(dispatcherProvider.main) { processNextCommand() }
-            } else {
-                Executors.MAIN_EXECUTOR.execute { processNextCommand() }
-            }
+            uiExecutor.execute { processNextCommand() }
         } else {
             Log.d(TAG, "not executed: $command - queue size: ${commandQueue.size}")
         }
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
index 5131774..de8be50 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
@@ -19,6 +19,7 @@
 import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static com.android.launcher3.MotionEventsUtils.isTrackpadScroll;
 import static com.android.launcher3.util.DisplayController.CHANGE_ALL;
 import static com.android.launcher3.util.DisplayController.CHANGE_NAVIGATION_MODE;
 import static com.android.launcher3.util.DisplayController.CHANGE_ROTATION;
@@ -70,7 +71,7 @@
 import com.android.launcher3.util.SettingsCache;
 import com.android.quickstep.TopTaskTracker.CachedTaskInfo;
 import com.android.quickstep.util.ActiveGestureLog;
-import com.android.quickstep.util.AssistStateManager;
+import com.android.quickstep.util.ContextualSearchStateManager;
 import com.android.quickstep.util.GestureExclusionManager;
 import com.android.quickstep.util.GestureExclusionManager.ExclusionListener;
 import com.android.quickstep.util.NavBarPosition;
@@ -101,7 +102,7 @@
     private final DisplayController mDisplayController;
 
     private final GestureExclusionManager mExclusionManager;
-    private final AssistStateManager mAssistStateManager;
+    private final ContextualSearchStateManager mContextualSearchStateManager;
 
     private final RotationTouchHelper mRotationTouchHelper;
     private final TaskStackChangeListener mPipListener;
@@ -152,7 +153,7 @@
         mContext = context;
         mDisplayController = DisplayController.INSTANCE.get(context);
         mExclusionManager = exclusionManager;
-        mAssistStateManager = AssistStateManager.INSTANCE.get(context);
+        mContextualSearchStateManager = ContextualSearchStateManager.INSTANCE.get(context);
         mIsOneHandedModeSupported = SystemProperties.getBoolean(SUPPORT_ONE_HANDED_MODE, false);
         mRotationTouchHelper = RotationTouchHelper.INSTANCE.get(context);
         if (isInstanceForTouches) {
@@ -563,6 +564,7 @@
         return mAssistantAvailable
                 && !QuickStepContract.isAssistantGestureDisabled(mSystemUiStateFlags)
                 && mRotationTouchHelper.touchInAssistantRegion(ev)
+                && !isTrackpadScroll(ev)
                 && !isLockToAppActive();
     }
 
@@ -617,8 +619,9 @@
                 : QUICKSTEP_TOUCH_SLOP_RATIO_TWO_BUTTON;
         float touchSlop = ViewConfiguration.get(mContext).getScaledTouchSlop();
 
-        if (mAssistStateManager.getLPNHCustomSlopMultiplier().isPresent()) {
-            float customSlopMultiplier = mAssistStateManager.getLPNHCustomSlopMultiplier().get();
+        if (mContextualSearchStateManager.getLPNHCustomSlopMultiplier().isPresent()) {
+            float customSlopMultiplier =
+                    mContextualSearchStateManager.getLPNHCustomSlopMultiplier().get();
             return customSlopMultiplier * slopMultiplier * touchSlop;
         } else {
             return slopMultiplier * touchSlop;
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index 5f02893..c66813f 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -64,7 +64,7 @@
 import com.android.launcher3.util.Preconditions;
 import com.android.launcher3.util.SafeCloseable;
 import com.android.quickstep.util.ActiveGestureProtoLogProxy;
-import com.android.quickstep.util.AssistUtils;
+import com.android.quickstep.util.ContextualSearchInvoker;
 import com.android.quickstep.util.unfold.ProxyUnfoldTransitionProvider;
 import com.android.systemui.shared.recents.ISystemUiProxy;
 import com.android.systemui.shared.recents.model.ThumbnailData;
@@ -311,8 +311,8 @@
         setBackToLauncherCallback(mBackToLauncherCallback, mBackToLauncherRunner);
         setUnfoldAnimationListener(mUnfoldAnimationListener);
         setDesktopTaskListener(mDesktopTaskListener);
-        setAssistantOverridesRequested(
-                AssistUtils.newInstance(mContext).getSysUiAssistOverrideInvocationTypes());
+        setAssistantOverridesRequested(ContextualSearchInvoker.newInstance(mContext)
+                .getSysUiAssistOverrideInvocationTypes());
         mStateChangeCallbacks.forEach(Runnable::run);
 
         if (mUnfoldTransitionProvider != null) {
@@ -1079,16 +1079,6 @@
         }
     }
 
-    public void removeFromSideStage(int taskId) {
-        if (mSplitScreen != null) {
-            try {
-                mSplitScreen.removeFromSideStage(taskId);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call removeFromSideStage");
-            }
-        }
-    }
-
     //
     // One handed
     //
@@ -1503,6 +1493,17 @@
         }
     }
 
+    /** Call shell to move a task with given `taskId` to external display. */
+    public void moveToExternalDisplay(int taskId) {
+        if (mDesktopMode != null) {
+            try {
+                mDesktopMode.moveToExternalDisplay(taskId);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call moveToExternalDisplay", e);
+            }
+        }
+    }
+
     //
     // Unfold transition
     //
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index 8e45767..0dbdcb7 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -116,6 +116,7 @@
             TaskShortcutFactory.INSTALL,
             TaskShortcutFactory.FREE_FORM,
             DesktopSystemShortcut.Companion.createFactory(),
+            ExternalDisplaySystemShortcut.Companion.createFactory(),
             TaskShortcutFactory.WELLBEING,
             TaskShortcutFactory.SAVE_APP_PAIR,
             TaskShortcutFactory.SCREENSHOT,
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index f0943dc..41a8a31 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -122,8 +122,8 @@
 import com.android.quickstep.util.ActiveGestureLog;
 import com.android.quickstep.util.ActiveGestureLog.CompoundString;
 import com.android.quickstep.util.ActiveGestureProtoLogProxy;
-import com.android.quickstep.util.AssistStateManager;
-import com.android.quickstep.util.AssistUtils;
+import com.android.quickstep.util.ContextualSearchInvoker;
+import com.android.quickstep.util.ContextualSearchStateManager;
 import com.android.quickstep.views.RecentsViewContainer;
 import com.android.systemui.shared.recents.IOverviewProxy;
 import com.android.systemui.shared.recents.ISystemUiProxy;
@@ -297,7 +297,8 @@
         @Override
         public void onAssistantOverrideInvoked(int invocationType) {
             executeForTouchInteractionService(tis -> {
-                if (!AssistUtils.newInstance(tis).tryStartAssistOverride(invocationType)) {
+                if (!ContextualSearchInvoker.newInstance(tis)
+                        .tryStartAssistOverride(invocationType)) {
                     Log.w(TAG, "Failed to invoke Assist override");
                 }
             });
@@ -1640,8 +1641,8 @@
         }
         mTaskbarManager.dumpLogs("", pw);
         mDesktopVisibilityController.dumpLogs("", pw);
-        pw.println("AssistStateManager:");
-        AssistStateManager.INSTANCE.get(this).dump("\t", pw);
+        pw.println("ContextualSearchStateManager:");
+        ContextualSearchStateManager.INSTANCE.get(this).dump("\t", pw);
         SystemUiProxy.INSTANCE.get(this).dump(pw);
         DeviceConfigWrapper.get().dump("   ", pw);
     }
diff --git a/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java b/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java
index 341c868..335161b 100644
--- a/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java
+++ b/quickstep/src/com/android/quickstep/dagger/QuickstepBaseAppComponent.java
@@ -21,6 +21,7 @@
 import com.android.launcher3.model.WellbeingModel;
 import com.android.quickstep.logging.SettingsChangeLogger;
 import com.android.quickstep.util.AsyncClockEventDelegate;
+import com.android.quickstep.util.ContextualSearchHapticManager;
 
 /**
  * Launcher Quickstep base component for Dagger injection.
@@ -36,4 +37,6 @@
     WellbeingModel getWellbeingModel();
 
     AsyncClockEventDelegate getAsyncClockEventDelegate();
+
+    ContextualSearchHapticManager getContextualSearchHapticManager();
 }
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
index e15fa54..fbf671f 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
@@ -84,10 +84,7 @@
  * [QuickstepProtoLogGroup.Constants.DEBUG_RECENTS_WINDOW]
  */
 class RecentsWindowManager(context: Context) :
-    RecentsWindowContext(context),
-    RecentsViewContainer,
-    StatefulContainer<RecentsState>,
-    RecentsAnimationListener {
+    RecentsWindowContext(context), RecentsViewContainer, StatefulContainer<RecentsState> {
 
     companion object {
         private const val HOME_APPEAR_DURATION: Long = 250
@@ -128,6 +125,17 @@
             }
         }
 
+    private val recentsAnimationListener =
+        object : RecentsAnimationListener {
+            override fun onRecentsAnimationCanceled(thumbnailDatas: HashMap<Int, ThumbnailData>) {
+                recentAnimationStopped()
+            }
+
+            override fun onRecentsAnimationFinished(controller: RecentsAnimationController) {
+                recentAnimationStopped()
+            }
+        }
+
     init {
         FallbackWindowInterface.init(this)
         TaskStackChangeListeners.getInstance().registerTaskStackListener(taskStackChangeListener)
@@ -138,7 +146,7 @@
         cleanupRecentsWindow()
         FallbackWindowInterface.getInstance()?.destroy()
         TaskStackChangeListeners.getInstance().unregisterTaskStackListener(taskStackChangeListener)
-        callbacks?.removeListener(this)
+        callbacks?.removeListener(recentsAnimationListener)
     }
 
     override fun startHome() {
@@ -203,7 +211,7 @@
             windowManager.removeViewImmediate(windowView)
         }
         stateManager.moveToRestState()
-        callbacks?.removeListener(this)
+        callbacks?.removeListener(recentsAnimationListener)
     }
 
     private fun isShowing(): Boolean {
@@ -249,17 +257,7 @@
         onInitListener?.test(true)
 
         this.callbacks = callbacks
-        callbacks?.addListener(this)
-    }
-
-    override fun onRecentsAnimationCanceled(thumbnailDatas: HashMap<Int, ThumbnailData>) {
-        super.onRecentsAnimationCanceled(thumbnailDatas)
-        recentAnimationStopped()
-    }
-
-    override fun onRecentsAnimationFinished(controller: RecentsAnimationController) {
-        super.onRecentsAnimationFinished(controller)
-        recentAnimationStopped()
+        callbacks?.addListener(recentsAnimationListener)
     }
 
     private fun recentAnimationStopped() {
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressHandler.java b/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressHandler.java
index 1d00e53..1a825a4 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressHandler.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressHandler.java
@@ -16,25 +16,67 @@
 
 package com.android.quickstep.inputconsumers;
 
+import static android.app.contextualsearch.ContextualSearchManager.ENTRYPOINT_LONG_PRESS_NAV_HANDLE;
+
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_ASSISTANT_SUCCESSFUL_NAV_HANDLE;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_OMNI_GET_LONG_PRESS_RUNNABLE;
+import static com.android.launcher3.logging.StatsLogManager.LauncherLatencyEvent.LAUNCHER_LATENCY_OMNI_RUNNABLE;
+
 import android.content.Context;
+import android.os.SystemClock;
+import android.util.Log;
+import android.view.ViewConfiguration;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
+import com.android.launcher3.LauncherApplication;
 import com.android.launcher3.R;
+import com.android.launcher3.logging.InstanceId;
+import com.android.launcher3.logging.InstanceIdSequence;
+import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.util.ResourceBasedOverride;
+import com.android.launcher3.util.VibratorWrapper;
+import com.android.quickstep.DeviceConfigWrapper;
 import com.android.quickstep.NavHandle;
+import com.android.quickstep.TopTaskTracker;
+import com.android.quickstep.util.ContextualSearchHapticManager;
+import com.android.quickstep.util.ContextualSearchInvoker;
+import com.android.quickstep.util.ContextualSearchStateManager;
 
 /**
  * Class for extending nav handle long press behavior
  */
 public class NavHandleLongPressHandler implements ResourceBasedOverride {
 
+    private static final String TAG = "NavHandleLongPressHandler";
+
+    protected final Context mContext;
+    protected final VibratorWrapper mVibratorWrapper;
+    protected final ContextualSearchHapticManager mContextualSearchHapticManager;
+    protected final ContextualSearchInvoker mContextualSearchInvoker;
+    protected final StatsLogManager mStatsLogManager;
+    private boolean mPendingInvocation;
+
+    public NavHandleLongPressHandler(Context context) {
+        mContext = context;
+        mStatsLogManager = StatsLogManager.newInstance(context);
+        mVibratorWrapper = VibratorWrapper.INSTANCE.get(mContext);
+        mContextualSearchHapticManager = ((LauncherApplication) context.getApplicationContext())
+                .getAppComponent().getContextualSearchHapticManager();
+        mContextualSearchInvoker = ContextualSearchInvoker.newInstance(mContext);
+    }
+
     /** Creates NavHandleLongPressHandler as specified by overrides */
     public static NavHandleLongPressHandler newInstance(Context context) {
         return Overrides.getObject(NavHandleLongPressHandler.class, context,
                 R.string.nav_handle_long_press_handler_class);
     }
 
+    protected boolean isContextualSearchEntrypointEnabled(NavHandle navHandle) {
+        return DeviceConfigWrapper.get().getEnableLongPressNavHandle();
+    }
+
     /**
      * Called when nav handle is long pressed to get the Runnable that should be executed by the
      * caller to invoke long press behavior. If null is returned that means long press couldn't be
@@ -46,8 +88,48 @@
      *
      * @param navHandle to handle this long press
      */
-    public @Nullable Runnable getLongPressRunnable(NavHandle navHandle) {
-        return null;
+    @Nullable
+    @VisibleForTesting
+    final Runnable getLongPressRunnable(NavHandle navHandle) {
+        if (!isContextualSearchEntrypointEnabled(navHandle)) {
+            Log.i(TAG, "Contextual Search invocation failed: entry point disabled");
+            mVibratorWrapper.cancelVibrate();
+            return null;
+        }
+
+        if (!mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures()) {
+            Log.i(TAG, "Contextual Search invocation failed: precondition not satisfied");
+            mVibratorWrapper.cancelVibrate();
+            return null;
+        }
+
+        mPendingInvocation = true;
+        Log.i(TAG, "Contextual Search invocation: invocation runnable created");
+        InstanceId instanceId = new InstanceIdSequence().newInstanceId();
+        mStatsLogManager.logger().withInstanceId(instanceId).log(
+                LAUNCHER_OMNI_GET_LONG_PRESS_RUNNABLE);
+        long startTimeMillis = SystemClock.elapsedRealtime();
+        return () -> {
+            mStatsLogManager.latencyLogger().withInstanceId(instanceId).withLatency(
+                    SystemClock.elapsedRealtime() - startTimeMillis).log(
+                    LAUNCHER_LATENCY_OMNI_RUNNABLE);
+            if (mContextualSearchInvoker.invokeContextualSearchUncheckedWithHaptic(
+                    ENTRYPOINT_LONG_PRESS_NAV_HANDLE)) {
+                Log.i(TAG, "Contextual Search invocation successful");
+
+                String runningPackage = TopTaskTracker.INSTANCE.get(mContext).getCachedTopTask(
+                        /* filterOnlyVisibleRecents */ true).getPackageName();
+                mStatsLogManager.logger().withPackageName(runningPackage)
+                        .log(LAUNCHER_LAUNCH_ASSISTANT_SUCCESSFUL_NAV_HANDLE);
+            } else {
+                mVibratorWrapper.cancelVibrate();
+                if (DeviceConfigWrapper.get().getAnimateLpnh()
+                        && !DeviceConfigWrapper.get().getShrinkNavHandleOnPress()) {
+                    navHandle.animateNavBarLongPress(
+                            /*isTouchDown*/false, /*shrink*/ false, /*durationMs*/160);
+                }
+            }
+        };
     }
 
     /**
@@ -55,7 +137,15 @@
      *
      * @param navHandle to handle the animation for this touch
      */
-    public void onTouchStarted(NavHandle navHandle) {}
+    @VisibleForTesting
+    final void onTouchStarted(NavHandle navHandle) {
+        mPendingInvocation = false;
+        if (isContextualSearchEntrypointEnabled(navHandle)
+                && mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures()) {
+            Log.i(TAG, "Contextual Search invocation: touch started");
+            startNavBarAnimation(navHandle);
+        }
+    }
 
     /**
      * Called when nav handle gesture is finished by the user lifting their finger or the system
@@ -64,5 +154,46 @@
      * @param navHandle to handle the animation for this touch
      * @param reason why the touch ended
      */
-    public void onTouchFinished(NavHandle navHandle, String reason) {}
+    @VisibleForTesting
+    final void onTouchFinished(NavHandle navHandle, String reason) {
+        Log.i(TAG, "Contextual Search invocation: touch finished with reason: " + reason);
+
+        if (!DeviceConfigWrapper.get().getShrinkNavHandleOnPress() || !mPendingInvocation) {
+            mVibratorWrapper.cancelVibrate();
+        }
+
+        if (DeviceConfigWrapper.get().getAnimateLpnh()) {
+            if (DeviceConfigWrapper.get().getShrinkNavHandleOnPress()) {
+                navHandle.animateNavBarLongPress(
+                        /*isTouchDown*/false, /*shrink*/ true, /*durationMs*/200);
+            } else {
+                navHandle.animateNavBarLongPress(
+                        /*isTouchDown*/false, /*shrink*/ false, /*durationMs*/ 160);
+            }
+        }
+    }
+
+    @VisibleForTesting
+    final void startNavBarAnimation(NavHandle navHandle) {
+        mContextualSearchHapticManager.vibrateForSearchHint();
+
+        if (DeviceConfigWrapper.get().getAnimateLpnh()) {
+            if (DeviceConfigWrapper.get().getShrinkNavHandleOnPress()) {
+                navHandle.animateNavBarLongPress(
+                        /*isTouchDown*/ true, /*shrink*/true, /*durationMs*/200);
+            } else {
+                long longPressTimeout;
+                ContextualSearchStateManager contextualSearchStateManager =
+                        ContextualSearchStateManager.INSTANCE.get(mContext);
+                if (contextualSearchStateManager.getLPNHDurationMillis().isPresent()) {
+                    longPressTimeout =
+                            contextualSearchStateManager.getLPNHDurationMillis().get().intValue();
+                } else {
+                    longPressTimeout = ViewConfiguration.getLongPressTimeout();
+                }
+                navHandle.animateNavBarLongPress(
+                        /*isTouchDown*/ true, /*shrink*/ false, /*durationMs*/ longPressTimeout);
+            }
+        }
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressInputConsumer.java
index f4d3695..f5bef05e 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressInputConsumer.java
@@ -38,7 +38,7 @@
 import com.android.quickstep.NavHandle;
 import com.android.quickstep.RecentsAnimationDeviceState;
 import com.android.quickstep.TopTaskTracker;
-import com.android.quickstep.util.AssistStateManager;
+import com.android.quickstep.util.ContextualSearchStateManager;
 import com.android.systemui.shared.system.InputMonitorCompat;
 
 /**
@@ -75,9 +75,11 @@
         super(delegate, inputMonitor);
         mScreenWidth = DisplayController.INSTANCE.get(context).getInfo().currentSize.x;
         mDeepPressEnabled = DeviceConfigWrapper.get().getEnableLpnhDeepPress();
-        AssistStateManager assistStateManager = AssistStateManager.INSTANCE.get(context);
-        if (assistStateManager.getLPNHDurationMillis().isPresent()) {
-            mLongPressTimeout = assistStateManager.getLPNHDurationMillis().get().intValue();
+        ContextualSearchStateManager contextualSearchStateManager =
+                ContextualSearchStateManager.INSTANCE.get(context);
+        if (contextualSearchStateManager.getLPNHDurationMillis().isPresent()) {
+            mLongPressTimeout =
+                    contextualSearchStateManager.getLPNHDurationMillis().get().intValue();
         } else {
             mLongPressTimeout = ViewConfiguration.getLongPressTimeout();
         }
diff --git a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
index b53650e..44b8b8d 100644
--- a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
+++ b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
@@ -34,6 +34,7 @@
 import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.quickstep.task.viewmodel.TaskOverlayViewModel
 import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
+import com.android.quickstep.task.viewmodel.TaskThumbnailViewModelImpl
 import com.android.quickstep.task.viewmodel.TaskViewData
 import com.android.quickstep.task.viewmodel.TaskViewModel
 import com.android.quickstep.views.TaskViewType
@@ -180,7 +181,7 @@
                 TaskContainerData::class.java -> TaskContainerData()
                 TaskThumbnailViewData::class.java -> TaskThumbnailViewData()
                 TaskThumbnailViewModel::class.java ->
-                    TaskThumbnailViewModel(
+                    TaskThumbnailViewModelImpl(
                         recentsViewData = inject(),
                         taskViewData = inject(scopeId, extras),
                         taskContainerData = inject(scopeId),
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
index 4970685..f55462a 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
@@ -10,144 +10,40 @@
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language goveryning permissions and
+ * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 
 package com.android.quickstep.task.viewmodel
 
-import android.annotation.ColorInt
-import android.app.ActivityTaskManager.INVALID_TASK_ID
 import android.graphics.Matrix
-import android.util.Log
-import androidx.core.graphics.ColorUtils
-import com.android.quickstep.recents.data.RecentTasksRepository
-import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
-import com.android.quickstep.recents.usecase.ThumbnailPositionState
-import com.android.quickstep.recents.viewmodel.RecentsViewData
-import com.android.quickstep.task.thumbnail.SplashAlphaUseCase
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.SnapshotSplash
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
-import com.android.systemui.shared.recents.model.Task
-import kotlin.math.max
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.flatMapLatest
-import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.flow.StateFlow
 
-@OptIn(ExperimentalCoroutinesApi::class)
-class TaskThumbnailViewModel(
-    recentsViewData: RecentsViewData,
-    taskViewData: TaskViewData,
-    taskContainerData: TaskContainerData,
-    private val tasksRepository: RecentTasksRepository,
-    private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase,
-    private val splashAlphaUseCase: SplashAlphaUseCase,
-) {
-    private val task = MutableStateFlow<Flow<Task?>>(flowOf(null))
-    private val splashProgress = MutableStateFlow(flowOf(0f))
-    private var taskId: Int = INVALID_TASK_ID
-
+/** ViewModel for representing TaskThumbnails */
+interface TaskThumbnailViewModel {
     /**
      * Progress for changes in corner radius. progress: 0 = overview corner radius; 1 = fullscreen
      * corner radius.
      */
-    val cornerRadiusProgress =
-        if (taskViewData.isOutlineFormedByThumbnailView) recentsViewData.fullscreenProgress
-        else MutableStateFlow(1f).asStateFlow()
+    val cornerRadiusProgress: StateFlow<Float>
 
-    val inheritedScale =
-        combine(recentsViewData.scale, taskViewData.scale) { recentsScale, taskScale ->
-            recentsScale * taskScale
-        }
+    /** The accumulated View.scale value for parent Views up to and including RecentsView */
+    val inheritedScale: Flow<Float>
 
-    val dimProgress: Flow<Float> =
-        combine(taskContainerData.taskMenuOpenProgress, recentsViewData.tintAmount) {
-            taskMenuOpenProgress,
-            tintAmount ->
-            max(taskMenuOpenProgress * MAX_SCRIM_ALPHA, tintAmount)
-        }
-    val splashAlpha = splashProgress.flatMapLatest { it }
+    /** Provides the level of dimming that the View should have */
+    val dimProgress: Flow<Float>
 
-    private val isLiveTile =
-        combine(
-                task.flatMapLatest { it }.map { it?.key?.id }.distinctUntilChanged(),
-                recentsViewData.runningTaskIds,
-                recentsViewData.runningTaskShowScreenshot,
-            ) { taskId, runningTaskIds, runningTaskShowScreenshot ->
-                runningTaskIds.contains(taskId) && !runningTaskShowScreenshot
-            }
-            .distinctUntilChanged()
+    /** Provides the alpha of the splash icon */
+    val splashAlpha: Flow<Float>
 
-    val uiState: Flow<TaskThumbnailUiState> =
-        combine(task.flatMapLatest { it }, isLiveTile) { taskVal, isRunning ->
-                // TODO(b/369339561) This log is firing a lot. Reduce emissions from TasksRepository
-                //  then re-enable this log.
-                //                Log.d(
-                //                    TAG,
-                //                    "Received task and / or live tile update. taskVal: $taskVal"
-                //                    + " isRunning: $isRunning.",
-                //                )
-                when {
-                    taskVal == null -> Uninitialized
-                    isRunning -> LiveTile
-                    isBackgroundOnly(taskVal) ->
-                        BackgroundOnly(taskVal.colorBackground.removeAlpha())
-                    isSnapshotSplashState(taskVal) ->
-                        SnapshotSplash(createSnapshotState(taskVal), taskVal.icon)
-                    else -> Uninitialized
-                }
-            }
-            .distinctUntilChanged()
+    /** Provides the UiState by which the task thumbnail can be represented */
+    val uiState: Flow<TaskThumbnailUiState>
 
-    fun bind(taskId: Int) {
-        Log.d(TAG, "bind taskId: $taskId")
-        this.taskId = taskId
-        task.value = tasksRepository.getTaskDataById(taskId)
-        splashProgress.value = splashAlphaUseCase.execute(taskId)
-    }
+    /** Attaches this ViewModel to a specific task id for it to provide data from. */
+    fun bind(taskId: Int)
 
-    fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix {
-        return runBlocking {
-            when (
-                val thumbnailPositionState =
-                    getThumbnailPositionUseCase.run(taskId, width, height, isRtl)
-            ) {
-                is ThumbnailPositionState.MatrixScaling -> thumbnailPositionState.matrix
-                is ThumbnailPositionState.MissingThumbnail -> Matrix.IDENTITY_MATRIX
-            }
-        }
-    }
-
-    private fun isBackgroundOnly(task: Task): Boolean = task.isLocked || task.thumbnail == null
-
-    private fun isSnapshotSplashState(task: Task): Boolean {
-        val thumbnailPresent = task.thumbnail?.thumbnail != null
-        val taskLocked = task.isLocked
-
-        return thumbnailPresent && !taskLocked
-    }
-
-    private fun createSnapshotState(task: Task): Snapshot {
-        val thumbnailData = task.thumbnail
-        val bitmap = thumbnailData?.thumbnail!!
-        return Snapshot(bitmap, thumbnailData.rotation, task.colorBackground.removeAlpha())
-    }
-
-    @ColorInt private fun Int.removeAlpha(): Int = ColorUtils.setAlphaComponent(this, 0xff)
-
-    private companion object {
-        const val MAX_SCRIM_ALPHA = 0.4f
-        const val TAG = "TaskThumbnailViewModel"
-    }
+    /** Returns a Matrix which can be applied to the snapshot */
+    fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix
 }
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
new file mode 100644
index 0000000..bd47cec
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
@@ -0,0 +1,149 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language goveryning permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.task.viewmodel
+
+import android.annotation.ColorInt
+import android.app.ActivityTaskManager.INVALID_TASK_ID
+import android.graphics.Matrix
+import android.util.Log
+import androidx.core.graphics.ColorUtils
+import com.android.quickstep.recents.data.RecentTasksRepository
+import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
+import com.android.quickstep.recents.usecase.ThumbnailPositionState
+import com.android.quickstep.recents.viewmodel.RecentsViewData
+import com.android.quickstep.task.thumbnail.SplashAlphaUseCase
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.SnapshotSplash
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
+import com.android.systemui.shared.recents.model.Task
+import kotlin.math.max
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.runBlocking
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class TaskThumbnailViewModelImpl(
+    recentsViewData: RecentsViewData,
+    taskViewData: TaskViewData,
+    taskContainerData: TaskContainerData,
+    private val tasksRepository: RecentTasksRepository,
+    private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase,
+    private val splashAlphaUseCase: SplashAlphaUseCase,
+) : TaskThumbnailViewModel {
+    private val task = MutableStateFlow<Flow<Task?>>(flowOf(null))
+    private val splashProgress = MutableStateFlow(flowOf(0f))
+    private var taskId: Int = INVALID_TASK_ID
+
+    override val cornerRadiusProgress =
+        if (taskViewData.isOutlineFormedByThumbnailView) recentsViewData.fullscreenProgress
+        else MutableStateFlow(1f).asStateFlow()
+
+    override val inheritedScale =
+        combine(recentsViewData.scale, taskViewData.scale) { recentsScale, taskScale ->
+            recentsScale * taskScale
+        }
+
+    override val dimProgress: Flow<Float> =
+        combine(taskContainerData.taskMenuOpenProgress, recentsViewData.tintAmount) {
+            taskMenuOpenProgress,
+            tintAmount ->
+            max(taskMenuOpenProgress * MAX_SCRIM_ALPHA, tintAmount)
+        }
+    override val splashAlpha = splashProgress.flatMapLatest { it }
+
+    private val isLiveTile =
+        combine(
+                task.flatMapLatest { it }.map { it?.key?.id }.distinctUntilChanged(),
+                recentsViewData.runningTaskIds,
+                recentsViewData.runningTaskShowScreenshot,
+            ) { taskId, runningTaskIds, runningTaskShowScreenshot ->
+                runningTaskIds.contains(taskId) && !runningTaskShowScreenshot
+            }
+            .distinctUntilChanged()
+
+    override val uiState: Flow<TaskThumbnailUiState> =
+        combine(task.flatMapLatest { it }, isLiveTile) { taskVal, isRunning ->
+                // TODO(b/369339561) This log is firing a lot. Reduce emissions from TasksRepository
+                //  then re-enable this log.
+                //                Log.d(
+                //                    TAG,
+                //                    "Received task and / or live tile update. taskVal: $taskVal"
+                //                    + " isRunning: $isRunning.",
+                //                )
+                when {
+                    taskVal == null -> Uninitialized
+                    isRunning -> LiveTile
+                    isBackgroundOnly(taskVal) ->
+                        BackgroundOnly(taskVal.colorBackground.removeAlpha())
+                    isSnapshotSplashState(taskVal) ->
+                        SnapshotSplash(createSnapshotState(taskVal), taskVal.icon)
+                    else -> Uninitialized
+                }
+            }
+            .distinctUntilChanged()
+
+    override fun bind(taskId: Int) {
+        Log.d(TAG, "bind taskId: $taskId")
+        this.taskId = taskId
+        task.value = tasksRepository.getTaskDataById(taskId)
+        splashProgress.value = splashAlphaUseCase.execute(taskId)
+    }
+
+    override fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix {
+        return runBlocking {
+            when (
+                val thumbnailPositionState =
+                    getThumbnailPositionUseCase.run(taskId, width, height, isRtl)
+            ) {
+                is ThumbnailPositionState.MatrixScaling -> thumbnailPositionState.matrix
+                is ThumbnailPositionState.MissingThumbnail -> Matrix.IDENTITY_MATRIX
+            }
+        }
+    }
+
+    private fun isBackgroundOnly(task: Task): Boolean = task.isLocked || task.thumbnail == null
+
+    private fun isSnapshotSplashState(task: Task): Boolean {
+        val thumbnailPresent = task.thumbnail?.thumbnail != null
+        val taskLocked = task.isLocked
+
+        return thumbnailPresent && !taskLocked
+    }
+
+    private fun createSnapshotState(task: Task): Snapshot {
+        val thumbnailData = task.thumbnail
+        val bitmap = thumbnailData?.thumbnail!!
+        return Snapshot(bitmap, thumbnailData.rotation, task.colorBackground.removeAlpha())
+    }
+
+    @ColorInt private fun Int.removeAlpha(): Int = ColorUtils.setAlphaComponent(this, 0xff)
+
+    private companion object {
+        const val MAX_SCRIM_ALPHA = 0.4f
+        const val TAG = "TaskThumbnailViewModel"
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/util/AssistStateManager.java b/quickstep/src/com/android/quickstep/util/AssistStateManager.java
deleted file mode 100644
index 7acb28d..0000000
--- a/quickstep/src/com/android/quickstep/util/AssistStateManager.java
+++ /dev/null
@@ -1,93 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.quickstep.util;
-
-import static com.android.launcher3.util.MainThreadInitializedObject.forOverride;
-
-import com.android.launcher3.R;
-import com.android.launcher3.util.MainThreadInitializedObject;
-import com.android.launcher3.util.ResourceBasedOverride;
-import com.android.launcher3.util.SafeCloseable;
-
-import java.io.PrintWriter;
-import java.util.Optional;
-
-/** Class to manage Assistant states. */
-public class AssistStateManager implements ResourceBasedOverride, SafeCloseable {
-
-    public static final MainThreadInitializedObject<AssistStateManager> INSTANCE =
-            forOverride(AssistStateManager.class, R.string.assist_state_manager_class);
-
-    public AssistStateManager() {}
-
-    /** Return {@code true} if the Settings toggle is enabled. */
-    public boolean isSettingsAllEntrypointsEnabled() {
-        return false;
-    }
-
-    /** Whether search supports showing on the lockscreen. */
-    public boolean supportsShowWhenLocked() {
-        return false;
-    }
-
-    /** Whether ContextualSearchService invocation path is available. */
-    public boolean isContextualSearchServiceAvailable() {
-        return false;
-    }
-
-    /** Get the Launcher overridden long press nav handle duration to trigger Assistant. */
-    public Optional<Long> getLPNHDurationMillis() {
-        return Optional.empty();
-    }
-
-    /**
-     * Get the Launcher overridden long press nav handle touch slop multiplier to trigger Assistant.
-     */
-    public Optional<Float> getLPNHCustomSlopMultiplier() {
-        return Optional.empty();
-    }
-
-    /** Get the Launcher overridden long press home duration to trigger Assistant. */
-    public Optional<Long> getLPHDurationMillis() {
-        return Optional.empty();
-    }
-
-    /** Get the Launcher overridden long press home touch slop multiplier to trigger Assistant. */
-    public Optional<Float> getLPHCustomSlopMultiplier() {
-        return Optional.empty();
-    }
-
-    /** Get the long press duration data source. */
-    public int getDurationDataSource() {
-        return 0;
-    }
-
-    /** Get the long press touch slop multiplier data source. */
-    public int getSlopDataSource() {
-        return 0;
-    }
-
-    /** Get the haptic bit overridden by AGSA. */
-    public Optional<Boolean> getShouldPlayHapticOverride() {
-        return Optional.empty();
-    }
-
-    /** Dump states. */
-    public void dump(String prefix, PrintWriter writer) {}
-
-    @Override
-    public void close() {}
-}
diff --git a/quickstep/src/com/android/quickstep/util/AssistUtils.java b/quickstep/src/com/android/quickstep/util/AssistUtils.java
deleted file mode 100644
index 11b6ea7..0000000
--- a/quickstep/src/com/android/quickstep/util/AssistUtils.java
+++ /dev/null
@@ -1,45 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.quickstep.util;
-
-import android.content.Context;
-
-import com.android.launcher3.R;
-import com.android.launcher3.util.ResourceBasedOverride;
-
-/** Utilities to work with Assistant functionality. */
-public class AssistUtils implements ResourceBasedOverride {
-
-    public AssistUtils() {}
-
-    /** Creates AssistUtils as specified by overrides */
-    public static AssistUtils newInstance(Context context) {
-        return Overrides.getObject(AssistUtils.class, context, R.string.assist_utils_class);
-    }
-
-    /** @return Array of AssistUtils.INVOCATION_TYPE_* that we want to handle instead of SysUI. */
-    public int[] getSysUiAssistOverrideInvocationTypes() {
-        return new int[0];
-    }
-
-    /**
-     * @return {@code true} if the override was handled, i.e. an assist surface was shown or the
-     * request should be ignored. {@code false} means the caller should start assist another way.
-     */
-    public boolean tryStartAssistOverride(int invocationType) {
-        return false;
-    }
-}
diff --git a/quickstep/src/com/android/quickstep/util/ContextualSearchHapticManager.kt b/quickstep/src/com/android/quickstep/util/ContextualSearchHapticManager.kt
new file mode 100644
index 0000000..8c246a5
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/util/ContextualSearchHapticManager.kt
@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.util
+
+import android.content.Context
+import android.os.VibrationEffect
+import android.os.VibrationEffect.Composition
+import android.os.Vibrator
+import com.android.launcher3.dagger.ApplicationContext
+import com.android.launcher3.dagger.LauncherAppSingleton
+import com.android.launcher3.util.VibratorWrapper
+import com.android.quickstep.DeviceConfigWrapper.Companion.get
+import javax.inject.Inject
+import kotlin.math.pow
+
+/** Manages haptics relating to Contextual Search invocations. */
+@LauncherAppSingleton
+class ContextualSearchHapticManager
+@Inject
+internal constructor(@ApplicationContext private val context: Context) {
+
+    private var searchEffect = createSearchEffect()
+    private var contextualSearchStateManager = ContextualSearchStateManager.INSTANCE[context]
+
+    private fun createSearchEffect() =
+        if (
+            context
+                .getSystemService(Vibrator::class.java)!!
+                .areAllPrimitivesSupported(Composition.PRIMITIVE_TICK)
+        ) {
+            VibrationEffect.startComposition()
+                .addPrimitive(Composition.PRIMITIVE_TICK, 1f)
+                .compose()
+        } else {
+            // fallback for devices without composition support
+            VibrationEffect.createPredefined(VibrationEffect.EFFECT_HEAVY_CLICK)
+        }
+
+    /** Indicates that search has been invoked. */
+    fun vibrateForSearch() {
+        searchEffect.let { VibratorWrapper.INSTANCE[context].vibrate(it) }
+    }
+
+    /** Indicates that search will be invoked if the current gesture is maintained. */
+    fun vibrateForSearchHint() {
+        val navbarConfig = get()
+        // Whether we should play the hint (ramp up) haptic
+        val shouldVibrate: Boolean =
+            if (
+                context
+                    .getSystemService(Vibrator::class.java)!!
+                    .areAllPrimitivesSupported(Composition.PRIMITIVE_LOW_TICK)
+            ) {
+                if (contextualSearchStateManager.shouldPlayHapticOverride.isPresent) {
+                    contextualSearchStateManager.shouldPlayHapticOverride.get()
+                } else {
+                    navbarConfig.enableSearchHapticHint
+                }
+            } else {
+                false
+            }
+
+        if (shouldVibrate) {
+            val startScale = navbarConfig.lpnhHapticHintStartScalePercent / 100f
+            val endScale = navbarConfig.lpnhHapticHintEndScalePercent / 100f
+            val scaleExponent = navbarConfig.lpnhHapticHintScaleExponent
+            val iterations = navbarConfig.lpnhHapticHintIterations
+            val delayMs = navbarConfig.lpnhHapticHintDelay
+            val composition = VibrationEffect.startComposition()
+            for (i in 0 until iterations) {
+                val t = i / (iterations - 1f)
+                val scale =
+                    ((1 - t) * startScale + t * endScale)
+                        .toDouble()
+                        .pow(scaleExponent.toDouble())
+                        .toFloat()
+                if (i == 0) {
+                    // Adds a delay before the ramp starts
+                    composition.addPrimitive(Composition.PRIMITIVE_LOW_TICK, scale, delayMs)
+                } else {
+                    composition.addPrimitive(Composition.PRIMITIVE_LOW_TICK, scale)
+                }
+            }
+            VibratorWrapper.INSTANCE[context].vibrate(composition.compose())
+        }
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/util/ContextualSearchInvoker.kt b/quickstep/src/com/android/quickstep/util/ContextualSearchInvoker.kt
new file mode 100644
index 0000000..dcb72aa
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/util/ContextualSearchInvoker.kt
@@ -0,0 +1,224 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.util
+
+import android.app.contextualsearch.ContextualSearchManager
+import android.app.contextualsearch.ContextualSearchManager.ENTRYPOINT_LONG_PRESS_HOME
+import android.app.contextualsearch.ContextualSearchManager.FEATURE_CONTEXTUAL_SEARCH
+import android.content.Context
+import android.util.Log
+import com.android.internal.app.AssistUtils
+import com.android.launcher3.LauncherApplication
+import com.android.launcher3.R
+import com.android.launcher3.logging.StatsLogManager
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_ASSISTANT_FAILED_SERVICE_ERROR
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_KEYGUARD
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_NOTIFICATION_SHADE
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_ATTEMPTED_SPLITSCREEN
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_FAILED_NOT_AVAILABLE
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_FAILED_SETTING_DISABLED
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_SUCCESSFUL_HOME
+import com.android.launcher3.util.ResourceBasedOverride
+import com.android.quickstep.DeviceConfigWrapper
+import com.android.quickstep.SystemUiProxy
+import com.android.quickstep.TopTaskTracker
+import com.android.systemui.shared.system.QuickStepContract
+
+/** Handles invocations and checks for Contextual Search. */
+open class ContextualSearchInvoker
+internal constructor(
+    protected val context: Context,
+    private val contextualSearchStateManager: ContextualSearchStateManager,
+    private val topTaskTracker: TopTaskTracker,
+    private val systemUiProxy: SystemUiProxy,
+    protected val statsLogManager: StatsLogManager,
+    private val contextualSearchHapticManager: ContextualSearchHapticManager,
+    private val contextualSearchManager: ContextualSearchManager?,
+) : ResourceBasedOverride {
+    constructor(
+        context: Context
+    ) : this(
+        context,
+        ContextualSearchStateManager.INSTANCE[context],
+        TopTaskTracker.INSTANCE[context],
+        SystemUiProxy.INSTANCE[context],
+        StatsLogManager.newInstance(context),
+        (context.applicationContext as LauncherApplication)
+            .appComponent
+            .contextualSearchHapticManager,
+        context.getSystemService(ContextualSearchManager::class.java),
+    )
+
+    /** @return Array of AssistUtils.INVOCATION_TYPE_* that we want to handle instead of SysUI. */
+    open fun getSysUiAssistOverrideInvocationTypes(): IntArray {
+        val overrideInvocationTypes = com.android.launcher3.util.IntArray()
+        if (context.packageManager.hasSystemFeature(FEATURE_CONTEXTUAL_SEARCH)) {
+            overrideInvocationTypes.add(AssistUtils.INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS)
+        }
+        return overrideInvocationTypes.toArray()
+    }
+
+    /**
+     * @return `true` if the override was handled, i.e. an assist surface was shown or the request
+     *   should be ignored. `false` means the caller should start assist another way.
+     */
+    fun tryStartAssistOverride(invocationType: Int): Boolean {
+        if (invocationType == AssistUtils.INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS) {
+            if (!context.packageManager.hasSystemFeature(FEATURE_CONTEXTUAL_SEARCH)) {
+                // When Contextual Search is disabled, fall back to Assistant.
+                return false
+            }
+
+            val success = show(ENTRYPOINT_LONG_PRESS_HOME)
+            if (success) {
+                val runningPackage =
+                    TopTaskTracker.INSTANCE[context].getCachedTopTask(
+                            /* filterOnlyVisibleRecents */ true
+                        )
+                        .getPackageName()
+                statsLogManager
+                    .logger()
+                    .withPackageName(runningPackage)
+                    .log(LAUNCHER_LAUNCH_OMNI_SUCCESSFUL_HOME)
+            }
+
+            // Regardless of success, do not fall back to other assistant.
+            return true
+        }
+        return false
+    }
+
+    /**
+     * Invoke Contextual Search via ContextualSearchService if availability checks are successful
+     *
+     * @param entryPoint one of the ENTRY_POINT_* constants defined in this class
+     * @return true if invocation was successful, false otherwise
+     */
+    fun show(entryPoint: Int): Boolean {
+        return if (!runContextualSearchInvocationChecksAndLogFailures()) false
+        else invokeContextualSearchUnchecked(entryPoint)
+    }
+
+    /**
+     * Run availability checks and log errors to WW. If successful the caller is expected to call
+     * {@link invokeContextualSearchUnchecked}
+     *
+     * @return true if availability checks were successful, false otherwise.
+     */
+    fun runContextualSearchInvocationChecksAndLogFailures(): Boolean {
+        if (
+            contextualSearchManager == null ||
+                !context.packageManager.hasSystemFeature(FEATURE_CONTEXTUAL_SEARCH)
+        ) {
+            Log.i(TAG, "Contextual Search invocation failed: no ContextualSearchManager")
+            statsLogManager.logger().log(LAUNCHER_LAUNCH_ASSISTANT_FAILED_SERVICE_ERROR)
+            return false
+        }
+        if (!contextualSearchStateManager.isContextualSearchSettingEnabled) {
+            Log.i(TAG, "Contextual Search invocation failed: setting disabled")
+            statsLogManager.logger().log(LAUNCHER_LAUNCH_OMNI_FAILED_SETTING_DISABLED)
+            return false
+        }
+        if (isNotificationShadeShowing()) {
+            Log.i(TAG, "Contextual Search invocation failed: notification shade")
+            statsLogManager.logger().log(LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_NOTIFICATION_SHADE)
+            return false
+        }
+        if (isKeyguardShowing()) {
+            Log.i(TAG, "Contextual Search invocation attempted: keyguard")
+            statsLogManager.logger().log(LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_KEYGUARD)
+            if (!contextualSearchStateManager.isInvocationAllowedOnKeyguard) {
+                Log.i(TAG, "Contextual Search invocation failed: keyguard not allowed")
+                return false
+            } else if (!contextualSearchStateManager.supportsShowWhenLocked()) {
+                Log.i(TAG, "Contextual Search invocation failed: AGA doesn't support keyguard")
+                return false
+            }
+        }
+        if (isInSplitscreen()) {
+            Log.i(TAG, "Contextual Search invocation attempted: splitscreen")
+            statsLogManager.logger().log(LAUNCHER_LAUNCH_OMNI_ATTEMPTED_SPLITSCREEN)
+            if (!contextualSearchStateManager.isInvocationAllowedInSplitscreen) {
+                Log.i(TAG, "Contextual Search invocation failed: splitscreen not allowed")
+                return false
+            }
+        }
+        if (!contextualSearchStateManager.isContextualSearchIntentAvailable) {
+            Log.i(TAG, "Contextual Search invocation failed: no matching CSS intent filter")
+            statsLogManager.logger().log(LAUNCHER_LAUNCH_OMNI_FAILED_NOT_AVAILABLE)
+            return false
+        }
+
+        return true
+    }
+
+    /**
+     * Invoke Contextual Search via ContextualSearchService and do haptic
+     *
+     * @param entryPoint Entry point identifier, passed to ContextualSearchService.
+     * @return true if invocation was successful, false otherwise
+     */
+    fun invokeContextualSearchUncheckedWithHaptic(entryPoint: Int): Boolean {
+        return invokeContextualSearchUnchecked(entryPoint, withHaptic = true)
+    }
+
+    private fun invokeContextualSearchUnchecked(
+        entryPoint: Int,
+        withHaptic: Boolean = false,
+    ): Boolean {
+        if (withHaptic && DeviceConfigWrapper.get().enableSearchHapticCommit) {
+            contextualSearchHapticManager.vibrateForSearch()
+        }
+        if (contextualSearchManager == null) {
+            return false
+        }
+        contextualSearchManager.startContextualSearch(entryPoint)
+        return true
+    }
+
+    private fun isInSplitscreen(): Boolean {
+        return topTaskTracker.getRunningSplitTaskIds().isNotEmpty()
+    }
+
+    private fun isNotificationShadeShowing(): Boolean {
+        return systemUiProxy.lastSystemUiStateFlags and SHADE_EXPANDED_SYSUI_FLAGS != 0L
+    }
+
+    private fun isKeyguardShowing(): Boolean {
+        return systemUiProxy.lastSystemUiStateFlags and KEYGUARD_SHOWING_SYSUI_FLAGS != 0L
+    }
+
+    companion object {
+        private const val TAG = "ContextualSearchInvoker"
+        const val SHADE_EXPANDED_SYSUI_FLAGS =
+            QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_EXPANDED or
+                QuickStepContract.SYSUI_STATE_QUICK_SETTINGS_EXPANDED
+        const val KEYGUARD_SHOWING_SYSUI_FLAGS =
+            (QuickStepContract.SYSUI_STATE_BOUNCER_SHOWING or
+                QuickStepContract.SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING or
+                QuickStepContract.SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING_OCCLUDED)
+
+        @JvmStatic
+        fun newInstance(context: Context): ContextualSearchInvoker {
+            return ResourceBasedOverride.Overrides.getObject(
+                ContextualSearchInvoker::class.java,
+                context,
+                R.string.contextual_search_invoker_class,
+            )
+        }
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/util/ContextualSearchStateManager.java b/quickstep/src/com/android/quickstep/util/ContextualSearchStateManager.java
new file mode 100644
index 0000000..142fc58
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/util/ContextualSearchStateManager.java
@@ -0,0 +1,281 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.quickstep.util;
+
+import static android.app.contextualsearch.ContextualSearchManager.ACTION_LAUNCH_CONTEXTUAL_SEARCH;
+import static android.app.contextualsearch.ContextualSearchManager.ENTRYPOINT_SYSTEM_ACTION;
+import static android.app.contextualsearch.ContextualSearchManager.FEATURE_CONTEXTUAL_SEARCH;
+
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_SUCCESSFUL_SYSTEM_ACTION;
+import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
+import static com.android.launcher3.util.MainThreadInitializedObject.forOverride;
+import static com.android.quickstep.util.SystemActionConstants.SYSTEM_ACTION_ID_SEARCH_SCREEN;
+
+import android.annotation.Nullable;
+import android.app.PendingIntent;
+import android.app.RemoteAction;
+import android.content.Context;
+import android.content.IIntentReceiver;
+import android.content.IIntentSender;
+import android.content.Intent;
+import android.content.pm.PackageManager;
+import android.graphics.drawable.Icon;
+import android.net.Uri;
+import android.os.Bundle;
+import android.os.IBinder;
+import android.os.RemoteException;
+import android.provider.Settings;
+import android.util.Log;
+import android.view.accessibility.AccessibilityManager;
+
+import androidx.annotation.CallSuper;
+import androidx.annotation.VisibleForTesting;
+
+import com.android.launcher3.R;
+import com.android.launcher3.logging.StatsLogManager;
+import com.android.launcher3.util.EventLogArray;
+import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.ResourceBasedOverride;
+import com.android.launcher3.util.SafeCloseable;
+import com.android.launcher3.util.SettingsCache;
+import com.android.launcher3.util.SimpleBroadcastReceiver;
+import com.android.quickstep.DeviceConfigWrapper;
+import com.android.quickstep.SystemUiProxy;
+import com.android.quickstep.TopTaskTracker;
+
+import java.io.PrintWriter;
+import java.util.Optional;
+
+/** Long-lived class to manage Contextual Search states like the user setting and availability. */
+public class ContextualSearchStateManager implements ResourceBasedOverride, SafeCloseable {
+
+    public static final MainThreadInitializedObject<ContextualSearchStateManager> INSTANCE =
+            forOverride(ContextualSearchStateManager.class,
+                    R.string.contextual_search_state_manager_class);
+
+    private static final String TAG = "ContextualSearchStMgr";
+    private static final int MAX_DEBUG_EVENT_SIZE = 20;
+    private static final Uri SEARCH_ALL_ENTRYPOINTS_ENABLED_URI =
+            Settings.Secure.getUriFor(Settings.Secure.SEARCH_ALL_ENTRYPOINTS_ENABLED);
+
+    private final Runnable mSysUiStateChangeListener = this::updateOverridesToSysUi;
+    private final SimpleBroadcastReceiver mContextualSearchPackageReceiver =
+            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, (unused) -> requestUpdateProperties());
+    private final SettingsCache.OnChangeListener mContextualSearchSettingChangedListener =
+            this::onContextualSearchSettingChanged;
+    protected final EventLogArray mEventLogArray = new EventLogArray(TAG, MAX_DEBUG_EVENT_SIZE);
+
+    @Nullable private SettingsCache mSettingsCache;
+    // Cached value whether the ContextualSearch intent filter matched any enabled components.
+    private boolean mIsContextualSearchIntentAvailable;
+    private boolean mIsContextualSearchSettingEnabled;
+
+    protected Context mContext;
+    protected String mContextualSearchPackage;
+
+    public ContextualSearchStateManager() {}
+
+    public ContextualSearchStateManager(Context context) {
+        mContext = context;
+        mContextualSearchPackage = mContext.getResources().getString(
+                com.android.internal.R.string.config_defaultContextualSearchPackageName);
+
+        if (areAllContextualSearchFlagsDisabled()
+                || !context.getPackageManager().hasSystemFeature(FEATURE_CONTEXTUAL_SEARCH)) {
+            // If we had previously registered a SystemAction which is no longer valid, we need to
+            // unregister it here.
+            unregisterSearchScreenSystemAction();
+            // Don't listen for stuff we aren't gonna use.
+            return;
+        }
+
+        requestUpdateProperties();
+        registerSearchScreenSystemAction();
+        mContextualSearchPackageReceiver.registerPkgActions(
+                context, mContextualSearchPackage, Intent.ACTION_PACKAGE_ADDED,
+                Intent.ACTION_PACKAGE_CHANGED, Intent.ACTION_PACKAGE_REMOVED);
+
+        mSettingsCache = SettingsCache.INSTANCE.get(context);
+        mSettingsCache.register(SEARCH_ALL_ENTRYPOINTS_ENABLED_URI,
+                mContextualSearchSettingChangedListener);
+        onContextualSearchSettingChanged(
+                mSettingsCache.getValue(SEARCH_ALL_ENTRYPOINTS_ENABLED_URI));
+        SystemUiProxy.INSTANCE.get(mContext).addOnStateChangeListener(mSysUiStateChangeListener);
+    }
+
+    /** Return {@code true} if the Settings toggle is enabled. */
+    public final boolean isContextualSearchSettingEnabled() {
+        return mIsContextualSearchSettingEnabled;
+    }
+
+    private void onContextualSearchSettingChanged(boolean isEnabled) {
+        mIsContextualSearchSettingEnabled = isEnabled;
+    }
+
+    /** Whether search supports showing on the lockscreen. */
+    protected boolean supportsShowWhenLocked() {
+        return false;
+    }
+
+    /** Whether ContextualSearchService invocation path is available. */
+    @VisibleForTesting
+    protected final boolean isContextualSearchIntentAvailable() {
+        return mIsContextualSearchIntentAvailable;
+    }
+
+    /** Get the Launcher overridden long press nav handle duration to trigger Assistant. */
+    public Optional<Long> getLPNHDurationMillis() {
+        return Optional.empty();
+    }
+
+    /**
+     * Get the Launcher overridden long press nav handle touch slop multiplier to trigger Assistant.
+     */
+    public Optional<Float> getLPNHCustomSlopMultiplier() {
+        return Optional.empty();
+    }
+
+    /** Get the Launcher overridden long press home duration to trigger Assistant. */
+    public Optional<Long> getLPHDurationMillis() {
+        return Optional.empty();
+    }
+
+    /** Get the Launcher overridden long press home touch slop multiplier to trigger Assistant. */
+    public Optional<Float> getLPHCustomSlopMultiplier() {
+        return Optional.empty();
+    }
+
+    /** Get the long press duration data source. */
+    public int getDurationDataSource() {
+        return 0;
+    }
+
+    /** Get the long press touch slop multiplier data source. */
+    public int getSlopDataSource() {
+        return 0;
+    }
+
+    /** Get the haptic bit overridden by AGSA. */
+    public Optional<Boolean> getShouldPlayHapticOverride() {
+        return Optional.empty();
+    }
+
+    protected boolean isInvocationAllowedOnKeyguard() {
+        return false;
+    }
+
+    protected boolean isInvocationAllowedInSplitscreen() {
+        return true;
+    }
+
+    @CallSuper
+    protected boolean areAllContextualSearchFlagsDisabled() {
+        return !DeviceConfigWrapper.get().getEnableLongPressNavHandle();
+    }
+
+    @CallSuper
+    protected void requestUpdateProperties() {
+        UI_HELPER_EXECUTOR.execute(() -> {
+            // Check that Contextual Search intent filters are enabled.
+            Intent csIntent = new Intent(ACTION_LAUNCH_CONTEXTUAL_SEARCH).setPackage(
+                    mContextualSearchPackage);
+            mIsContextualSearchIntentAvailable =
+                    !mContext.getPackageManager().queryIntentActivities(csIntent,
+                            PackageManager.MATCH_DIRECT_BOOT_AWARE
+                                    | PackageManager.MATCH_DIRECT_BOOT_UNAWARE).isEmpty();
+
+            addEventLog("Updated isContextualSearchIntentAvailable",
+                    mIsContextualSearchIntentAvailable);
+        });
+    }
+
+    protected final void updateOverridesToSysUi() {
+        // LPH commit haptic is always enabled
+        SystemUiProxy.INSTANCE.get(mContext).setOverrideHomeButtonLongPress(
+                getLPHDurationMillis().orElse(0L), getLPHCustomSlopMultiplier().orElse(0f), true);
+        Log.i(TAG, "Sent LPH override to sysui: " + getLPHDurationMillis().orElse(0L) + ";"
+                + getLPHCustomSlopMultiplier().orElse(0f));
+    }
+
+    private void registerSearchScreenSystemAction() {
+        PendingIntent searchScreenPendingIntent = new PendingIntent(new IIntentSender.Stub() {
+            @Override
+            public void send(int i, Intent intent, String s, IBinder iBinder,
+                    IIntentReceiver iIntentReceiver, String s1, Bundle bundle)
+                    throws RemoteException {
+                // Delayed slightly to minimize chance of capturing the System Actions dialog.
+                UI_HELPER_EXECUTOR.getHandler().postDelayed(
+                        () -> {
+                            boolean contextualSearchInvoked =
+                                    ContextualSearchInvoker.newInstance(mContext).show(
+                                            ENTRYPOINT_SYSTEM_ACTION);
+                            if (contextualSearchInvoked) {
+                                String runningPackage =
+                                        TopTaskTracker.INSTANCE.get(mContext).getCachedTopTask(
+                                                /* filterOnlyVisibleRecents */
+                                                true).getPackageName();
+                                StatsLogManager.newInstance(mContext).logger()
+                                        .withPackageName(runningPackage)
+                                        .log(LAUNCHER_LAUNCH_OMNI_SUCCESSFUL_SYSTEM_ACTION);
+                            }
+                        }, 200);
+            }
+        });
+
+        mContext.getSystemService(AccessibilityManager.class).registerSystemAction(new RemoteAction(
+                        Icon.createWithResource(mContext, R.drawable.ic_allapps_search),
+                        mContext.getString(R.string.search_gesture_feature_title),
+                        mContext.getString(R.string.search_gesture_feature_title),
+                        searchScreenPendingIntent),
+                SYSTEM_ACTION_ID_SEARCH_SCREEN);
+    }
+
+    private void unregisterSearchScreenSystemAction() {
+        mContext.getSystemService(AccessibilityManager.class).unregisterSystemAction(
+                SYSTEM_ACTION_ID_SEARCH_SCREEN);
+    }
+
+    /** Dump states. */
+    public final void dump(String prefix, PrintWriter writer) {
+        synchronized (mEventLogArray) {
+            mEventLogArray.dump(prefix, writer);
+        }
+    }
+
+    @Override
+    public void close() {
+        mContextualSearchPackageReceiver.unregisterReceiverSafely(mContext);
+        unregisterSearchScreenSystemAction();
+
+        if (mSettingsCache != null) {
+            mSettingsCache.unregister(SEARCH_ALL_ENTRYPOINTS_ENABLED_URI,
+                    mContextualSearchSettingChangedListener);
+        }
+        SystemUiProxy.INSTANCE.get(mContext).removeOnStateChangeListener(mSysUiStateChangeListener);
+    }
+
+    protected final void addEventLog(String event) {
+        synchronized (mEventLogArray) {
+            mEventLogArray.addLog(event);
+        }
+    }
+
+    protected final void addEventLog(String event, boolean extras) {
+        synchronized (mEventLogArray) {
+            mEventLogArray.addLog(event, extras);
+        }
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
index 9335e7e..a5be89a 100644
--- a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
+++ b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
@@ -30,7 +30,6 @@
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
 import android.content.Context;
-import android.content.SharedPreferences;
 import android.graphics.Matrix;
 import android.graphics.Point;
 import android.graphics.PointF;
@@ -45,6 +44,7 @@
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherPrefChangeListener;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.testing.shared.TestProtocol;
 import com.android.launcher3.touch.PagedOrientationHandler;
@@ -66,8 +66,7 @@
  * This class has initial default state assuming the device and foreground app have
  * no ({@link Surface#ROTATION_0} rotation.
  */
-public class RecentsOrientedState implements
-        SharedPreferences.OnSharedPreferenceChangeListener {
+public class RecentsOrientedState implements LauncherPrefChangeListener {
 
     private static final String TAG = "RecentsOrientedState";
     private static final boolean DEBUG = false;
@@ -283,7 +282,7 @@
     }
 
     @Override
-    public void onSharedPreferenceChanged(SharedPreferences sharedPreferences, String s) {
+    public void onPrefChanged(String s) {
         if (LauncherPrefs.ALLOW_ROTATION.getSharedPrefKey().equals(s)) {
             updateHomeRotationSetting();
         }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 7554c44..bccfa0c 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -6673,6 +6673,26 @@
         successCallback.run();
     }
 
+    /**
+     * Move the provided task into external display and invoke {@code successCallback} if succeeded.
+     */
+    public void moveTaskToExternalDisplay(TaskContainer taskContainer, Runnable successCallback) {
+        if (!DesktopModeStatus.canEnterDesktopMode(mContext)) {
+            return;
+        }
+        switchToScreenshot(() -> finishRecentsAnimation(/* toRecents= */true, /* shouldPip= */false,
+                () -> moveTaskToDesktopInternal(taskContainer, successCallback)));
+    }
+
+    private void moveTaskToDesktopInternal(TaskContainer taskContainer, Runnable successCallback) {
+        if (mDesktopRecentsTransitionController == null) {
+            return;
+        }
+        mDesktopRecentsTransitionController.moveToExternalDisplay(taskContainer.getTask().key.id);
+        successCallback.run();
+    }
+
+
     // Logs when the orientation of Overview changes. We log both real and fake orientation changes.
     private void logOrientationChanged() {
         // Only log when Overview is showing.
diff --git a/quickstep/src_protolog/com/android/launcher3/util/StateManagerProtoLogProxy.java b/quickstep/src_protolog/com/android/launcher3/util/StateManagerProtoLogProxy.java
new file mode 100644
index 0000000..bc989dc
--- /dev/null
+++ b/quickstep/src_protolog/com/android/launcher3/util/StateManagerProtoLogProxy.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util;
+
+import static com.android.launcher3.Flags.enableStateManagerProtoLog;
+import static com.android.quickstep.util.QuickstepProtoLogGroup.LAUNCHER_STATE_MANAGER;
+
+import androidx.annotation.NonNull;
+
+import com.android.internal.protolog.ProtoLog;
+
+/**
+ * Proxy class used for StateManager ProtoLog support.
+ */
+public class StateManagerProtoLogProxy {
+
+    public static void logGoToState(
+            @NonNull Object fromState, @NonNull Object toState, @NonNull String trace) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER,
+                "StateManager.goToState: fromState: %s, toState: %s, partial trace:\n%s",
+                fromState,
+                toState,
+                trace);
+    }
+
+    public static void logCreateAtomicAnimation(
+            @NonNull Object fromState, @NonNull Object toState, @NonNull String trace) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER, "StateManager.createAtomicAnimation: "
+                        + "fromState: %s, toState: %s, partial trace:\n%s",
+                fromState,
+                toState,
+                trace);
+    }
+
+    public static void logOnStateTransitionStart(@NonNull Object state) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER, "StateManager.onStateTransitionStart: state: %s", state);
+    }
+
+    public static void logOnStateTransitionEnd(@NonNull Object state) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER, "StateManager.onStateTransitionEnd: state: %s", state);
+    }
+
+    public static void logCancelAnimation(boolean animationOngoing, @NonNull String trace) {
+        if (!enableStateManagerProtoLog()) return;
+        ProtoLog.d(LAUNCHER_STATE_MANAGER,
+                "StateManager.cancelAnimation: animation ongoing: %b, partial trace:\n%s",
+                animationOngoing,
+                trace);
+    }
+}
diff --git a/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java b/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java
index 7b81b9a..bb02a11 100644
--- a/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java
+++ b/quickstep/src_protolog/com/android/quickstep/util/QuickstepProtoLogGroup.java
@@ -27,7 +27,8 @@
 public enum QuickstepProtoLogGroup implements IProtoLogGroup {
 
     ACTIVE_GESTURE_LOG(true, true, false, "ActiveGestureLog"),
-    RECENTS_WINDOW(true, true, Constants.DEBUG_RECENTS_WINDOW, "RecentsWindow");
+    RECENTS_WINDOW(true, true, Constants.DEBUG_RECENTS_WINDOW, "RecentsWindow"),
+    LAUNCHER_STATE_MANAGER(true, true, Constants.DEBUG_STATE_MANAGER, "LauncherStateManager");
 
     private final boolean mEnabled;
     private volatile boolean mLogToProto;
@@ -97,6 +98,7 @@
     private static final class Constants {
 
         private static final boolean DEBUG_RECENTS_WINDOW = false;
+        private static final boolean DEBUG_STATE_MANAGER = true; // b/279059025, b/325463989
 
         private static final int LOG_START_ID =
                 (int) (UUID.nameUUIDFromBytes(QuickstepProtoLogGroup.class.getName().getBytes())
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt
new file mode 100644
index 0000000..ff5d8bd
--- /dev/null
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.task.thumbnail
+
+import android.graphics.Matrix
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
+import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
+import kotlinx.coroutines.flow.MutableStateFlow
+
+class FakeTaskThumbnailViewModel : TaskThumbnailViewModel {
+    override val cornerRadiusProgress = MutableStateFlow(0f)
+    override val inheritedScale = MutableStateFlow(1f)
+    override val dimProgress = MutableStateFlow(0f)
+    override val splashAlpha = MutableStateFlow(0f)
+    override val uiState = MutableStateFlow<TaskThumbnailUiState>(Uninitialized)
+
+    override fun bind(taskId: Int) {
+        // no-op
+    }
+
+    override fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean) =
+        Matrix.IDENTITY_MATRIX
+}
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt
new file mode 100644
index 0000000..75769e9
--- /dev/null
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt
@@ -0,0 +1,86 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.quickstep.task.thumbnail
+
+import android.content.Context
+import android.graphics.Color
+import android.view.LayoutInflater
+import com.android.launcher3.R
+import com.android.quickstep.recents.di.RecentsDependencies
+import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
+import com.google.android.apps.nexuslauncher.imagecomparison.goldenpathmanager.ViewScreenshotGoldenPathManager
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
+import platform.test.screenshot.DeviceEmulationSpec
+import platform.test.screenshot.Displays
+import platform.test.screenshot.ViewScreenshotTestRule
+import platform.test.screenshot.getEmulatedDevicePathConfig
+
+/** Screenshot tests for [TaskThumbnailView]. */
+@RunWith(ParameterizedAndroidJunit4::class)
+class TaskThumbnailViewScreenshotTest(emulationSpec: DeviceEmulationSpec) {
+
+    @get:Rule
+    val screenshotRule =
+        ViewScreenshotTestRule(
+            emulationSpec,
+            ViewScreenshotGoldenPathManager(getEmulatedDevicePathConfig(emulationSpec)),
+        )
+
+    private val taskThumbnailViewModel = FakeTaskThumbnailViewModel()
+
+    @Test
+    fun taskThumbnailView_uninitialized() {
+        screenshotRule.screenshotTest("taskThumbnailView_uninitialized") { activity ->
+            activity.actionBar?.hide()
+            createTaskThumbnailView(activity)
+        }
+    }
+
+    @Test
+    fun taskThumbnailView_backgroundOnly() {
+        screenshotRule.screenshotTest("taskThumbnailView_backgroundOnly") { activity ->
+            activity.actionBar?.hide()
+            taskThumbnailViewModel.uiState.value = TaskThumbnailUiState.BackgroundOnly(Color.YELLOW)
+            createTaskThumbnailView(activity)
+        }
+    }
+
+    private fun createTaskThumbnailView(context: Context): TaskThumbnailView {
+        val di = RecentsDependencies.initialize(context)
+        val taskThumbnailView =
+            LayoutInflater.from(context).inflate(R.layout.task_thumbnail, null, false)
+        val ttvDiScopeId = di.getScope(taskThumbnailView).scopeId
+        di.provide(TaskThumbnailViewData::class.java, ttvDiScopeId) { TaskThumbnailViewData() }
+        di.provide(TaskThumbnailViewModel::class.java, ttvDiScopeId) { taskThumbnailViewModel }
+
+        return taskThumbnailView as TaskThumbnailView
+    }
+
+    companion object {
+        @Parameters(name = "{0}")
+        @JvmStatic
+        fun getTestSpecs() =
+            DeviceEmulationSpec.forDisplays(
+                Displays.Phone,
+                isDarkTheme = false,
+                isLandscape = false,
+            )
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarControllerTestUtil.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarControllerTestUtil.kt
index a57fb70..6e2f74a 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarControllerTestUtil.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarControllerTestUtil.kt
@@ -16,10 +16,34 @@
 
 package com.android.launcher3.taskbar
 
+import android.content.Context
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.ConstantItem
+import com.android.launcher3.LauncherPrefs
+import kotlin.properties.ReadWriteProperty
+import kotlin.reflect.KProperty
 
 object TaskbarControllerTestUtil {
     inline fun runOnMainSync(crossinline runTest: () -> Unit) {
         getInstrumentation().runOnMainSync { runTest() }
     }
+
+    /** Returns a property to read/write the value of a [ConstantItem]. */
+    fun <T : Any> ConstantItem<T>.asProperty(context: Context): ReadWriteProperty<Any?, T> {
+        return TaskbarItemProperty(context, this)
+    }
+
+    private class TaskbarItemProperty<T : Any>(
+        private val context: Context,
+        private val item: ConstantItem<T>,
+    ) : ReadWriteProperty<Any?, T> {
+
+        override fun getValue(thisRef: Any?, property: KProperty<*>): T {
+            return LauncherPrefs.get(context).get(item)
+        }
+
+        override fun setValue(thisRef: Any?, property: KProperty<*>, value: T) {
+            runOnMainSync { LauncherPrefs.get(context).put(item, value) }
+        }
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarEduTooltipControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarEduTooltipControllerTest.kt
index e575efd..e60717b 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarEduTooltipControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarEduTooltipControllerTest.kt
@@ -14,24 +14,17 @@
  * limitations under the License.
  */
 
-package com.android.launcher3.taskbar.test
+package com.android.launcher3.taskbar
 
 import android.util.Log
 import com.android.launcher3.Utilities
-import com.android.launcher3.taskbar.TOOLTIP_STEP_FEATURES
-import com.android.launcher3.taskbar.TOOLTIP_STEP_NONE
-import com.android.launcher3.taskbar.TOOLTIP_STEP_PINNING
-import com.android.launcher3.taskbar.TOOLTIP_STEP_SWIPE
-import com.android.launcher3.taskbar.TaskbarActivityContext
+import com.android.launcher3.taskbar.TaskbarControllerTestUtil.asProperty
 import com.android.launcher3.taskbar.TaskbarControllerTestUtil.runOnMainSync
-import com.android.launcher3.taskbar.TaskbarEduTooltipController
 import com.android.launcher3.taskbar.rules.TaskbarModeRule
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.PINNED
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.THREE_BUTTONS
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.TRANSIENT
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
-import com.android.launcher3.taskbar.rules.TaskbarPinningPreferenceRule
-import com.android.launcher3.taskbar.rules.TaskbarPreferenceRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
@@ -53,19 +46,9 @@
 
     @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
 
-    @get:Rule(order = 1)
-    val tooltipStepPreferenceRule =
-        TaskbarPreferenceRule(context, OnboardingPrefs.TASKBAR_EDU_TOOLTIP_STEP.prefItem)
+    @get:Rule(order = 1) val taskbarModeRule = TaskbarModeRule(context)
 
-    @get:Rule(order = 2)
-    val searchEduPreferenceRule =
-        TaskbarPreferenceRule(context, OnboardingPrefs.TASKBAR_SEARCH_EDU_SEEN)
-
-    @get:Rule(order = 3) val taskbarPinningPreferenceRule = TaskbarPinningPreferenceRule(context)
-
-    @get:Rule(order = 4) val taskbarModeRule = TaskbarModeRule(context)
-
-    @get:Rule(order = 5) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 2) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @InjectController lateinit var taskbarEduTooltipController: TaskbarEduTooltipController
 
@@ -74,6 +57,9 @@
 
     private val wasInTestHarness = Utilities.isRunningInTestHarness()
 
+    private var tooltipStep by OnboardingPrefs.TASKBAR_EDU_TOOLTIP_STEP.prefItem.asProperty(context)
+    private var searchEduSeen by OnboardingPrefs.TASKBAR_SEARCH_EDU_SEEN.asProperty(context)
+
     @Before
     fun setUp() {
         Log.e("Taskbar", "TaskbarEduTooltipControllerTest test started")
@@ -91,7 +77,7 @@
     @Test
     @TaskbarMode(THREE_BUTTONS)
     fun testMaybeShowSwipeEdu_whenTaskbarIsInThreeButtonMode_doesNotShowSwipeEdu() {
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_SWIPE
+        tooltipStep = TOOLTIP_STEP_SWIPE
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_SWIPE)
         runOnMainSync { taskbarEduTooltipController.maybeShowSwipeEdu() }
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_SWIPE)
@@ -101,7 +87,7 @@
     @Test
     @TaskbarMode(TRANSIENT)
     fun testMaybeShowSwipeEdu_whenSwipeEduAlreadyShown_doesNotShowSwipeEdu() {
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_FEATURES
+        tooltipStep = TOOLTIP_STEP_FEATURES
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_FEATURES)
         runOnMainSync { taskbarEduTooltipController.maybeShowSwipeEdu() }
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_FEATURES)
@@ -111,7 +97,7 @@
     @Test
     @TaskbarMode(TRANSIENT)
     fun testMaybeShowSwipeEdu_whenUserHasNotSeen_doesShowSwipeEdu() {
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_SWIPE
+        tooltipStep = TOOLTIP_STEP_SWIPE
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_SWIPE)
         runOnMainSync { taskbarEduTooltipController.maybeShowSwipeEdu() }
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_FEATURES)
@@ -121,7 +107,7 @@
     @Test
     @TaskbarMode(TRANSIENT)
     fun testMaybeShowFeaturesEdu_whenFeatureEduAlreadyShown_doesNotShowFeatureEdu() {
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_NONE
+        tooltipStep = TOOLTIP_STEP_NONE
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_NONE)
         runOnMainSync { taskbarEduTooltipController.maybeShowFeaturesEdu() }
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_NONE)
@@ -131,7 +117,7 @@
     @Test
     @TaskbarMode(TRANSIENT)
     fun testMaybeShowFeaturesEdu_whenUserHasNotSeen_doesShowFeatureEdu() {
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_FEATURES
+        tooltipStep = TOOLTIP_STEP_FEATURES
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_FEATURES)
         runOnMainSync { taskbarEduTooltipController.maybeShowFeaturesEdu() }
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_NONE)
@@ -141,7 +127,7 @@
     @Test
     @TaskbarMode(THREE_BUTTONS)
     fun testMaybeShowPinningEdu_whenTaskbarIsInThreeButtonMode_doesNotShowPinningEdu() {
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_PINNING
+        tooltipStep = TOOLTIP_STEP_PINNING
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_PINNING)
         runOnMainSync { taskbarEduTooltipController.maybeShowFeaturesEdu() }
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_PINNING)
@@ -152,7 +138,7 @@
     @TaskbarMode(TRANSIENT)
     fun testMaybeShowPinningEdu_whenUserHasNotSeen_doesShowPinningEdu() {
         // Test standalone pinning edu, where user has seen taskbar edu before, but not pinning edu.
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_PINNING
+        tooltipStep = TOOLTIP_STEP_PINNING
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_PINNING)
         runOnMainSync { taskbarEduTooltipController.maybeShowFeaturesEdu() }
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_NONE)
@@ -162,21 +148,21 @@
     @Test
     @TaskbarMode(TRANSIENT)
     fun testIsBeforeTooltipFeaturesStep_whenUserHasNotSeenFeatureEdu_shouldReturnTrue() {
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_SWIPE
+        tooltipStep = TOOLTIP_STEP_SWIPE
         assertThat(taskbarEduTooltipController.isBeforeTooltipFeaturesStep).isTrue()
     }
 
     @Test
     @TaskbarMode(TRANSIENT)
     fun testIsBeforeTooltipFeaturesStep_whenUserHasSeenFeatureEdu_shouldReturnFalse() {
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_NONE
+        tooltipStep = TOOLTIP_STEP_NONE
         assertThat(taskbarEduTooltipController.isBeforeTooltipFeaturesStep).isFalse()
     }
 
     @Test
     @TaskbarMode(TRANSIENT)
     fun testHide_whenTooltipIsOpen_shouldCloseTooltip() {
-        tooltipStepPreferenceRule.value = TOOLTIP_STEP_SWIPE
+        tooltipStep = TOOLTIP_STEP_SWIPE
         assertThat(taskbarEduTooltipController.tooltipStep).isEqualTo(TOOLTIP_STEP_SWIPE)
         assertThat(taskbarEduTooltipController.isTooltipOpen).isFalse()
         runOnMainSync { taskbarEduTooltipController.maybeShowSwipeEdu() }
@@ -196,7 +182,7 @@
     @Test
     @TaskbarMode(PINNED)
     fun testMaybeShowSearchEdu_whenTaskbarIsPinnedAndUserHasSeenSearchEdu_shouldNotShowSearchEdu() {
-        searchEduPreferenceRule.value = true
+        searchEduSeen = true
         assertThat(taskbarEduTooltipController.userHasSeenSearchEdu).isTrue()
         runOnMainSync { taskbarEduTooltipController.hide() }
         assertThat(taskbarEduTooltipController.isTooltipOpen).isFalse()
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
index 02d6218..253d921 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
@@ -39,7 +39,7 @@
 import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarNavButtonCallbacks;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TouchInteractionService;
-import com.android.quickstep.util.AssistUtils;
+import com.android.quickstep.util.ContextualSearchInvoker;
 import com.android.systemui.contextualeducation.GestureType;
 
 import org.junit.Before;
@@ -64,7 +64,7 @@
     @Mock
     Handler mockHandler;
     @Mock
-    AssistUtils mockAssistUtils;
+    ContextualSearchInvoker mockContextualSearchInvoker;
     @Mock
     StatsLogManager mockStatsLogManager;
     @Mock
@@ -109,7 +109,7 @@
                 mockSystemUiProxy,
                 mockContextualEduStatsManager,
                 mockHandler,
-                mockAssistUtils);
+                mockContextualSearchInvoker);
     }
 
     @Test
@@ -166,40 +166,40 @@
     @Test
     public void testLongPressHome_enabled_withoutOverride() {
         mNavButtonController.setAssistantLongPressEnabled(true /*assistantLongPressEnabled*/);
-        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(false);
+        when(mockContextualSearchInvoker.tryStartAssistOverride(anyInt())).thenReturn(false);
 
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
-        verify(mockAssistUtils, times(1)).tryStartAssistOverride(anyInt());
+        verify(mockContextualSearchInvoker, times(1)).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, times(1)).startAssistant(any());
     }
 
     @Test
     public void testLongPressHome_enabled_withOverride() {
         mNavButtonController.setAssistantLongPressEnabled(true /*assistantLongPressEnabled*/);
-        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(true);
+        when(mockContextualSearchInvoker.tryStartAssistOverride(anyInt())).thenReturn(true);
 
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
-        verify(mockAssistUtils, times(1)).tryStartAssistOverride(anyInt());
+        verify(mockContextualSearchInvoker, times(1)).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, never()).startAssistant(any());
     }
 
     @Test
     public void testLongPressHome_disabled_withoutOverride() {
         mNavButtonController.setAssistantLongPressEnabled(false /*assistantLongPressEnabled*/);
-        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(false);
+        when(mockContextualSearchInvoker.tryStartAssistOverride(anyInt())).thenReturn(false);
 
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
-        verify(mockAssistUtils, never()).tryStartAssistOverride(anyInt());
+        verify(mockContextualSearchInvoker, never()).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, never()).startAssistant(any());
     }
 
     @Test
     public void testLongPressHome_disabled_withOverride() {
         mNavButtonController.setAssistantLongPressEnabled(false /*assistantLongPressEnabled*/);
-        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(true);
+        when(mockContextualSearchInvoker.tryStartAssistOverride(anyInt())).thenReturn(true);
 
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
-        verify(mockAssistUtils, never()).tryStartAssistOverride(anyInt());
+        verify(mockContextualSearchInvoker, never()).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, never()).startAssistant(any());
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
index de73ce7..71f4ef4 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
@@ -18,10 +18,13 @@
 
 import android.animation.AnimatorTestRule
 import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.LauncherPrefs.Companion.TASKBAR_PINNING
 import com.android.launcher3.R
 import com.android.launcher3.taskbar.StashedHandleViewController.ALPHA_INDEX_STASHED
 import com.android.launcher3.taskbar.TaskbarAutohideSuspendController.FLAG_AUTOHIDE_SUSPEND_EDU_OPEN
+import com.android.launcher3.taskbar.TaskbarControllerTestUtil.asProperty
 import com.android.launcher3.taskbar.TaskbarStashController.FLAG_IN_APP
 import com.android.launcher3.taskbar.TaskbarStashController.FLAG_IN_OVERVIEW
 import com.android.launcher3.taskbar.TaskbarStashController.FLAG_IN_STASHED_LAUNCHER_STATE
@@ -42,7 +45,6 @@
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.THREE_BUTTONS
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.TRANSIENT
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
-import com.android.launcher3.taskbar.rules.TaskbarPinningPreferenceRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.UserSetupMode
@@ -63,11 +65,11 @@
 @EnableFlags(FLAG_ENABLE_BUBBLE_BAR)
 @EmulatedDevices(["pixelTablet2023"])
 class TaskbarStashControllerTest {
-    @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
-    @get:Rule(order = 1) val taskbarModeRule = TaskbarModeRule(context)
-    @get:Rule(order = 2) val taskbarPinningPreferenceRule = TaskbarPinningPreferenceRule(context)
-    @get:Rule(order = 3) val animatorTestRule = AnimatorTestRule(this)
-    @get:Rule(order = 4) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 0) val setFlagsRule = SetFlagsRule()
+    @get:Rule(order = 1) val context = TaskbarWindowSandboxContext.create()
+    @get:Rule(order = 2) val taskbarModeRule = TaskbarModeRule(context)
+    @get:Rule(order = 4) val animatorTestRule = AnimatorTestRule(this)
+    @get:Rule(order = 5) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @InjectController lateinit var stashController: TaskbarStashController
     @InjectController lateinit var viewController: TaskbarViewController
@@ -121,10 +123,11 @@
 
     @Test
     fun testRecreateAsTransient_timeoutStarted() {
-        taskbarPinningPreferenceRule.isPinned = true
+        var isPinned by TASKBAR_PINNING.asProperty(context)
+        isPinned = true
         activityContext.controllers.sharedState?.taskbarWasPinned = true
 
-        taskbarPinningPreferenceRule.isPinned = false
+        isPinned = false
         assertThat(stashController.timeoutAlarm.alarmPending()).isTrue()
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewControllerTest.kt
index 516220a..3c0d9c6 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewControllerTest.kt
@@ -21,6 +21,7 @@
 import com.android.launcher3.appprediction.AppsDividerView
 import com.android.launcher3.appprediction.AppsDividerView.DividerType
 import com.android.launcher3.appprediction.PredictionRowView
+import com.android.launcher3.taskbar.TaskbarControllerTestUtil.asProperty
 import com.android.launcher3.taskbar.TaskbarStashController
 import com.android.launcher3.taskbar.TaskbarStashController.FLAG_STASHED_IN_APP_AUTO
 import com.android.launcher3.taskbar.allapps.TaskbarAllAppsControllerTest.Companion.TEST_PREDICTED_APPS
@@ -29,7 +30,6 @@
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.PINNED
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.TRANSIENT
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
-import com.android.launcher3.taskbar.rules.TaskbarPreferenceRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
@@ -49,14 +49,12 @@
 
     @get:Rule(order = 0) val context = TaskbarWindowSandboxContext.create()
     @get:Rule(order = 1) val taskbarModeRule = TaskbarModeRule(context)
-    @get:Rule(order = 2)
-    val allAppsVisitedPreferenceRule =
-        TaskbarPreferenceRule(context, ALL_APPS_VISITED_COUNT.prefItem)
-    @get:Rule(order = 3) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
+    @get:Rule(order = 2) val taskbarUnitTestRule = TaskbarUnitTestRule(this, context)
 
     @InjectController lateinit var overlayController: TaskbarOverlayController
     @InjectController lateinit var stashController: TaskbarStashController
 
+    private var allAppsVisitedCount by ALL_APPS_VISITED_COUNT.prefItem.asProperty(context)
     private val searchSessionController =
         TestUtil.getOnUiThread { TaskbarSearchSessionController.newInstance(context) }
 
@@ -102,7 +100,7 @@
 
     @Test
     fun testShow_firstAllAppsVisit_hasAllAppsTextDivider() {
-        allAppsVisitedPreferenceRule.value = 0
+        allAppsVisitedCount = 0
         val viewController = createViewController()
         getInstrumentation().runOnMainSync { viewController.show(false) }
 
@@ -120,7 +118,7 @@
 
     @Test
     fun testShow_maxAllAppsVisitedCount_hasLineDivider() {
-        allAppsVisitedPreferenceRule.value = ALL_APPS_VISITED_COUNT.maxCount
+        allAppsVisitedCount = ALL_APPS_VISITED_COUNT.maxCount
         val viewController = createViewController()
         getInstrumentation().runOnMainSync { viewController.show(false) }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
index fdafce0..3dd7689 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
@@ -22,12 +22,15 @@
 import android.view.Gravity
 import android.widget.FrameLayout
 import android.widget.TextView
+import androidx.core.animation.AnimatorTestRule
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
 import com.android.launcher3.R
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -36,11 +39,15 @@
 @RunWith(AndroidJUnit4::class)
 class BubbleBarFlyoutControllerTest {
 
+    @get:Rule val animatorTestRule = AnimatorTestRule()
+
     private lateinit var flyoutController: BubbleBarFlyoutController
     private lateinit var flyoutContainer: FrameLayout
+    private lateinit var topBoundaryListener: FakeTopBoundaryListener
     private val context = ApplicationProvider.getApplicationContext<Context>()
     private val flyoutMessage = BubbleBarFlyoutMessage(icon = null, "sender name", "message")
     private var onLeft = true
+    private var flyoutTy = 50f
 
     @Before
     fun setUp() {
@@ -50,53 +57,126 @@
                 override val isOnLeft
                     get() = onLeft
 
-                override val targetTy = 50f
+                override val targetTy
+                    get() = flyoutTy
+
                 override val distanceToCollapsedPosition = PointF(100f, 200f)
                 override val collapsedSize = 30f
                 override val collapsedColor = Color.BLUE
                 override val collapsedElevation = 1f
                 override val distanceToRevealTriangle = 50f
             }
-        flyoutController = BubbleBarFlyoutController(flyoutContainer, positioner)
+        topBoundaryListener = FakeTopBoundaryListener()
+        val flyoutScheduler = FlyoutScheduler { block -> block.invoke() }
+        flyoutController =
+            BubbleBarFlyoutController(
+                flyoutContainer,
+                positioner,
+                topBoundaryListener,
+                flyoutScheduler,
+            )
     }
 
     @Test
     fun flyoutPosition_left() {
-        flyoutController.setUpFlyout(flyoutMessage)
-        assertThat(flyoutContainer.childCount).isEqualTo(1)
-        val flyout = flyoutContainer.getChildAt(0)
-        val lp = flyout.layoutParams as FrameLayout.LayoutParams
-        assertThat(lp.gravity).isEqualTo(Gravity.BOTTOM or Gravity.LEFT)
-        assertThat(flyout.translationY).isEqualTo(50f)
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            flyoutController.setUpFlyout(flyoutMessage)
+            assertThat(flyoutContainer.childCount).isEqualTo(1)
+            val flyout = flyoutContainer.getChildAt(0)
+            val lp = flyout.layoutParams as FrameLayout.LayoutParams
+            assertThat(lp.gravity).isEqualTo(Gravity.BOTTOM or Gravity.LEFT)
+            assertThat(flyout.translationY).isEqualTo(50f)
+        }
     }
 
     @Test
     fun flyoutPosition_right() {
         onLeft = false
-        flyoutController.setUpFlyout(flyoutMessage)
-        assertThat(flyoutContainer.childCount).isEqualTo(1)
-        val flyout = flyoutContainer.getChildAt(0)
-        val lp = flyout.layoutParams as FrameLayout.LayoutParams
-        assertThat(lp.gravity).isEqualTo(Gravity.BOTTOM or Gravity.RIGHT)
-        assertThat(flyout.translationY).isEqualTo(50f)
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            flyoutController.setUpFlyout(flyoutMessage)
+            assertThat(flyoutContainer.childCount).isEqualTo(1)
+            val flyout = flyoutContainer.getChildAt(0)
+            val lp = flyout.layoutParams as FrameLayout.LayoutParams
+            assertThat(lp.gravity).isEqualTo(Gravity.BOTTOM or Gravity.RIGHT)
+            assertThat(flyout.translationY).isEqualTo(50f)
+        }
     }
 
     @Test
     fun flyoutMessage() {
-        flyoutController.setUpFlyout(flyoutMessage)
-        assertThat(flyoutContainer.childCount).isEqualTo(1)
-        val flyout = flyoutContainer.getChildAt(0)
-        val sender = flyout.findViewById<TextView>(R.id.bubble_flyout_title)
-        assertThat(sender.text).isEqualTo("sender name")
-        val message = flyout.findViewById<TextView>(R.id.bubble_flyout_text)
-        assertThat(message.text).isEqualTo("message")
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            flyoutController.setUpFlyout(flyoutMessage)
+            assertThat(flyoutContainer.childCount).isEqualTo(1)
+            val flyout = flyoutContainer.getChildAt(0)
+            val sender = flyout.findViewById<TextView>(R.id.bubble_flyout_title)
+            assertThat(sender.text).isEqualTo("sender name")
+            val message = flyout.findViewById<TextView>(R.id.bubble_flyout_text)
+            assertThat(message.text).isEqualTo("message")
+        }
     }
 
     @Test
     fun hideFlyout_removedFromContainer() {
-        flyoutController.setUpFlyout(flyoutMessage)
-        assertThat(flyoutContainer.childCount).isEqualTo(1)
-        flyoutController.hideFlyout()
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            flyoutController.setUpFlyout(flyoutMessage)
+            assertThat(flyoutContainer.childCount).isEqualTo(1)
+            flyoutController.hideFlyout {}
+            animatorTestRule.advanceTimeBy(300)
+        }
         assertThat(flyoutContainer.childCount).isEqualTo(0)
     }
+
+    @Test
+    fun showFlyout_extendsTopBoundary() {
+        // set negative translation for the flyout so that it will request to extend the top
+        // boundary
+        flyoutTy = -50f
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            flyoutController.setUpFlyout(flyoutMessage)
+            assertThat(flyoutContainer.childCount).isEqualTo(1)
+        }
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animatorTestRule.advanceTimeBy(300)
+        }
+        assertThat(topBoundaryListener.topBoundaryExtendedSpace).isEqualTo(50)
+    }
+
+    @Test
+    fun showFlyout_withinBoundary() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            flyoutController.setUpFlyout(flyoutMessage)
+            assertThat(flyoutContainer.childCount).isEqualTo(1)
+        }
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            animatorTestRule.advanceTimeBy(300)
+        }
+        assertThat(topBoundaryListener.topBoundaryExtendedSpace).isEqualTo(0)
+    }
+
+    @Test
+    fun hideFlyout_resetsTopBoundary() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            flyoutController.setUpFlyout(flyoutMessage)
+            assertThat(flyoutContainer.childCount).isEqualTo(1)
+            flyoutController.hideFlyout {}
+            animatorTestRule.advanceTimeBy(300)
+        }
+        assertThat(topBoundaryListener.topBoundaryReset).isTrue()
+    }
+
+    class FakeTopBoundaryListener : BubbleBarFlyoutController.TopBoundaryListener {
+
+        var topBoundaryExtendedSpace = 0
+        var topBoundaryReset = false
+
+        override fun extendTopBoundary(space: Int) {
+            topBoundaryExtendedSpace = space
+        }
+
+        override fun resetTopBoundary() {
+            topBoundaryReset = true
+        }
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRule.kt
deleted file mode 100644
index d417790..0000000
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRule.kt
+++ /dev/null
@@ -1,65 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.taskbar.rules
-
-import android.platform.test.flag.junit.FlagsParameterization
-import android.platform.test.flag.junit.SetFlagsRule
-import com.android.launcher3.Flags.FLAG_ENABLE_TASKBAR_PINNING
-import com.android.launcher3.LauncherPrefs.Companion.TASKBAR_PINNING
-import com.android.launcher3.LauncherPrefs.Companion.TASKBAR_PINNING_IN_DESKTOP_MODE
-import com.android.launcher3.util.DisplayController
-import org.junit.rules.RuleChain
-import org.junit.rules.TestRule
-import org.junit.runner.Description
-import org.junit.runners.model.Statement
-
-/**
- * Rule that allows modifying the Taskbar pinned preferences.
- *
- * The original preference values are restored on teardown.
- *
- * If this rule is being used with [TaskbarUnitTestRule], make sure this rule is applied first.
- *
- * This rule is overkill if a test does not need to change the mode during Taskbar's lifecycle. If
- * the mode is static, use [TaskbarModeRule] instead, which forces the mode. A test can class can
- * declare both this rule and [TaskbarModeRule] but using both for a test method is unsupported.
- */
-class TaskbarPinningPreferenceRule(context: TaskbarWindowSandboxContext) : TestRule {
-
-    private val setFlagsRule =
-        SetFlagsRule(FlagsParameterization(mapOf(FLAG_ENABLE_TASKBAR_PINNING to true)))
-    private val pinningRule = TaskbarPreferenceRule(context, TASKBAR_PINNING)
-    private val desktopPinningRule = TaskbarPreferenceRule(context, TASKBAR_PINNING_IN_DESKTOP_MODE)
-    private val ruleChain =
-        RuleChain.outerRule(setFlagsRule).around(pinningRule).around(desktopPinningRule)
-
-    var isPinned by pinningRule::value
-    var isPinnedInDesktopMode by desktopPinningRule::value
-
-    override fun apply(base: Statement, description: Description): Statement {
-        return object : Statement() {
-            override fun evaluate() {
-                DisplayController.enableTaskbarModePreferenceForTests(true)
-                try {
-                    ruleChain.apply(base, description).evaluate()
-                } finally {
-                    DisplayController.enableTaskbarModePreferenceForTests(false)
-                }
-            }
-        }
-    }
-}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt
deleted file mode 100644
index 977e7a5..0000000
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPinningPreferenceRuleTest.kt
+++ /dev/null
@@ -1,114 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.taskbar.rules
-
-import com.android.launcher3.util.DisplayController
-import com.android.launcher3.util.LauncherMultivalentJUnit
-import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
-import com.android.launcher3.util.window.WindowManagerProxy
-import com.google.android.apps.nexuslauncher.deviceemulator.TestWindowManagerProxy
-import com.google.common.truth.Truth.assertThat
-import org.junit.Rule
-import org.junit.Test
-import org.junit.runner.Description
-import org.junit.runner.RunWith
-import org.junit.runners.model.Statement
-
-@RunWith(LauncherMultivalentJUnit::class)
-@EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
-class TaskbarPinningPreferenceRuleTest {
-    @get:Rule val context = TaskbarWindowSandboxContext.create()
-
-    private val preferenceRule = TaskbarPinningPreferenceRule(context)
-
-    @Test
-    fun testEnablePinning_verifyDisplayController() {
-        onSetup {
-            preferenceRule.isPinned = true
-            preferenceRule.isPinnedInDesktopMode = false
-            assertThat(DisplayController.isPinnedTaskbar(context)).isTrue()
-        }
-    }
-
-    @Test
-    fun testDisablePinning_verifyDisplayController() {
-        onSetup {
-            preferenceRule.isPinned = false
-            preferenceRule.isPinnedInDesktopMode = false
-            assertThat(DisplayController.isPinnedTaskbar(context)).isFalse()
-        }
-    }
-
-    @Test
-    fun testEnableDesktopPinning_verifyDisplayController() {
-        context.putObject(
-            WindowManagerProxy.INSTANCE,
-            TestWindowManagerProxy(context).apply { isInDesktopMode = true },
-        )
-
-        onSetup {
-            preferenceRule.isPinned = false
-            preferenceRule.isPinnedInDesktopMode = true
-            assertThat(DisplayController.isPinnedTaskbar(context)).isTrue()
-        }
-    }
-
-    @Test
-    fun testDisableDesktopPinning_verifyDisplayController() {
-        context.putObject(
-            WindowManagerProxy.INSTANCE,
-            TestWindowManagerProxy(context).apply { isInDesktopMode = true },
-        )
-
-        onSetup {
-            preferenceRule.isPinned = false
-            preferenceRule.isPinnedInDesktopMode = false
-            assertThat(DisplayController.isPinnedTaskbar(context)).isFalse()
-        }
-    }
-
-    @Test
-    fun testTearDown_afterTogglingPinnedPreference_preferenceReset() {
-        val wasPinned = preferenceRule.isPinned
-        onSetup { preferenceRule.isPinned = !preferenceRule.isPinned }
-        assertThat(preferenceRule.isPinned).isEqualTo(wasPinned)
-    }
-
-    @Test
-    fun testTearDown_afterTogglingDesktopPreference_preferenceReset() {
-        val wasPinnedInDesktopMode = preferenceRule.isPinnedInDesktopMode
-        onSetup { preferenceRule.isPinnedInDesktopMode = !preferenceRule.isPinnedInDesktopMode }
-        assertThat(preferenceRule.isPinnedInDesktopMode).isEqualTo(wasPinnedInDesktopMode)
-    }
-
-    /** Executes [runTest] after the [preferenceRule] setup phase completes. */
-    private fun onSetup(runTest: () -> Unit) {
-        preferenceRule
-            .apply(
-                object : Statement() {
-                    override fun evaluate() = runTest()
-                },
-                DESCRIPTION,
-            )
-            .evaluate()
-    }
-
-    private companion object {
-        private val DESCRIPTION =
-            Description.createSuiteDescription(TaskbarPinningPreferenceRule::class.java)
-    }
-}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRule.kt
deleted file mode 100644
index e42ca9e..0000000
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRule.kt
+++ /dev/null
@@ -1,55 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.taskbar.rules
-
-import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
-import com.android.launcher3.ConstantItem
-import com.android.launcher3.LauncherPrefs
-import org.junit.rules.TestRule
-import org.junit.runner.Description
-import org.junit.runners.model.Statement
-
-/**
- * Rule for modifying a Taskbar preference.
- *
- * The original preference value is restored on teardown.
- */
-class TaskbarPreferenceRule<T : Any>(
-    private val context: TaskbarWindowSandboxContext,
-    private val constantItem: ConstantItem<T>,
-) : TestRule {
-
-    private val prefs: LauncherPrefs
-        get() = LauncherPrefs.get(context)
-
-    var value: T
-        get() = prefs.get(constantItem)
-        set(value) = getInstrumentation().runOnMainSync { prefs.put(constantItem, value) }
-
-    override fun apply(base: Statement, description: Description): Statement {
-        return object : Statement() {
-            override fun evaluate() {
-                val originalValue = value
-                try {
-                    base.evaluate()
-                } finally {
-                    value = originalValue
-                }
-            }
-        }
-    }
-}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRuleTest.kt
deleted file mode 100644
index b7e6fa3..0000000
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarPreferenceRuleTest.kt
+++ /dev/null
@@ -1,67 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.taskbar.rules
-
-import com.android.launcher3.LauncherPrefs.Companion.TASKBAR_PINNING
-import com.android.launcher3.util.LauncherMultivalentJUnit
-import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
-import com.google.common.truth.Truth.assertThat
-import org.junit.Rule
-import org.junit.Test
-import org.junit.runner.Description
-import org.junit.runner.RunWith
-import org.junit.runners.model.Statement
-
-@RunWith(LauncherMultivalentJUnit::class)
-@EmulatedDevices(["pixelFoldable2023", "pixelTablet2023"])
-class TaskbarPreferenceRuleTest {
-
-    @get:Rule val context = TaskbarWindowSandboxContext.create()
-    private val preferenceRule = TaskbarPreferenceRule(context, TASKBAR_PINNING)
-
-    @Test
-    fun testSetup_toggleBoolean_updatesPreferences() {
-        val originalValue = preferenceRule.value
-        onSetup {
-            preferenceRule.value = !preferenceRule.value
-            assertThat(preferenceRule.value).isNotEqualTo(originalValue)
-        }
-    }
-
-    @Test
-    fun testTeardown_afterTogglingBoolean_preferenceReset() {
-        val originalValue = preferenceRule.value
-        onSetup { preferenceRule.value = !preferenceRule.value }
-        assertThat(preferenceRule.value).isEqualTo(originalValue)
-    }
-
-    private fun onSetup(runTest: () -> Unit) {
-        preferenceRule
-            .apply(
-                object : Statement() {
-                    override fun evaluate() = runTest()
-                },
-                DESCRIPTION,
-            )
-            .evaluate()
-    }
-
-    private companion object {
-        private val DESCRIPTION =
-            Description.createSuiteDescription(TaskbarPreferenceRule::class.java)
-    }
-}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
index 741be50..2d3bfd6 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
@@ -16,19 +16,67 @@
 
 package com.android.launcher3.taskbar.rules
 
+import android.content.Context
 import android.content.ContextWrapper
+import android.hardware.display.DisplayManager
+import android.hardware.display.VirtualDisplay
+import android.view.Display.DEFAULT_DISPLAY
+import androidx.test.core.app.ApplicationProvider
+import com.android.launcher3.FakeLauncherPrefs
+import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.util.MainThreadInitializedObject.ObjectSandbox
 import com.android.launcher3.util.SandboxApplication
+import org.junit.rules.ExternalResource
+import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
 
-/** Sandbox Context for running Taskbar tests. */
-class TaskbarWindowSandboxContext private constructor(base: SandboxApplication) :
-    ContextWrapper(base), ObjectSandbox by base, TestRule by base {
+/**
+ * [SandboxApplication] for running Taskbar tests.
+ *
+ * Tests need to run on a [VirtualDisplay] to avoid conflicting with Launcher's Taskbar on the
+ * [DEFAULT_DISPLAY] (i.e. test is executing on a device).
+ */
+class TaskbarWindowSandboxContext
+private constructor(base: SandboxApplication, val virtualDisplay: VirtualDisplay) :
+    ContextWrapper(base),
+    ObjectSandbox by base,
+    TestRule by RuleChain.outerRule(virtualDisplayRule(virtualDisplay)).around(base) {
+
+    init {
+        putObject(LauncherPrefs.INSTANCE, FakeLauncherPrefs(this))
+    }
 
     companion object {
+        private const val VIRTUAL_DISPLAY_NAME = "TaskbarSandboxDisplay"
+
         /** Creates a [SandboxApplication] for Taskbar tests. */
         fun create(): TaskbarWindowSandboxContext {
-            return TaskbarWindowSandboxContext(SandboxApplication())
+            val base = ApplicationProvider.getApplicationContext<Context>()
+            val displayManager = checkNotNull(base.getSystemService(DisplayManager::class.java))
+
+            // Create virtual display to avoid clashing with Taskbar on default display.
+            val virtualDisplay =
+                base.resources.displayMetrics.let {
+                    displayManager.createVirtualDisplay(
+                        VIRTUAL_DISPLAY_NAME,
+                        it.widthPixels,
+                        it.heightPixels,
+                        it.densityDpi,
+                        /* surface= */ null,
+                        /* flags= */ 0,
+                    )
+                }
+
+            return TaskbarWindowSandboxContext(
+                SandboxApplication(base.createDisplayContext(virtualDisplay.display)),
+                virtualDisplay,
+            )
         }
     }
 }
+
+private fun virtualDisplayRule(virtualDisplay: VirtualDisplay): TestRule {
+    return object : ExternalResource() {
+        override fun after() = virtualDisplay.release()
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
new file mode 100644
index 0000000..69095e7
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.taskbar.rules
+
+import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.Description
+import org.junit.runner.RunWith
+import org.junit.runners.model.Statement
+
+@RunWith(LauncherMultivalentJUnit::class)
+@EmulatedDevices(["pixelFoldable2023"])
+class TaskbarWindowSandboxContextTest {
+
+    @Test
+    fun testVirtualDisplay_releasedOnTeardown() {
+        val context = TaskbarWindowSandboxContext.create()
+        assertThat(context.virtualDisplay.token).isNotNull()
+
+        context
+            .apply(
+                object : Statement() {
+                    override fun evaluate() = Unit
+                },
+                Description.createSuiteDescription(TaskbarWindowSandboxContextTest::class.java),
+            )
+            .evaluate()
+
+        assertThat(context.virtualDisplay.token).isNull()
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
index 0ae710f..b0db737 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
@@ -41,7 +41,6 @@
 import org.junit.runner.RunWith
 import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.spy
-import org.mockito.Mockito.`when`
 import org.mockito.kotlin.any
 import org.mockito.kotlin.mock
 
@@ -56,10 +55,12 @@
     private val testScope = TestScope(dispatcher)
 
     private var pendingCallbacksWithDelays = mutableListOf<Long>()
+    private lateinit var pendingCommandsToExecute: MutableList<Runnable>
 
     @Suppress("UNCHECKED_CAST")
     @Before
     fun setup() {
+        pendingCommandsToExecute = mutableListOf()
         setFlagsRule.setFlags(true, Flags.FLAG_ENABLE_OVERVIEW_COMMAND_HELPER_TIMEOUT)
 
         sut =
@@ -68,7 +69,8 @@
                     touchInteractionService = mock(),
                     overviewComponentObserver = mock(),
                     taskAnimationManager = mock(),
-                    dispatcherProvider = TestDispatcherProvider(dispatcher)
+                    dispatcherProvider = TestDispatcherProvider(dispatcher),
+                    uiExecutor = { runnable -> pendingCommandsToExecute += runnable },
                 )
             )
 
@@ -94,12 +96,21 @@
         pendingCallbacksWithDelays.add(delayInMillis)
     }
 
+    /**
+     * This function runs all the pending commands from the Executor for testing purposes. Replacing
+     * the uiExecutor allows the test to execute the command queue manually, making it possible to
+     * assert each state of the commands in the queue individually.
+     */
+    private fun executePendingCommands() = pendingCommandsToExecute.forEach { it.run() }
+
     @Test
     fun whenFirstCommandIsAdded_executeCommandImmediately() =
         testScope.runTest {
             // Add command to queue
             val commandInfo: CommandInfo = sut.addCommand(CommandType.HOME)!!
             assertThat(commandInfo.status).isEqualTo(CommandStatus.IDLE)
+            executePendingCommands()
+            assertThat(commandInfo.status).isEqualTo(CommandStatus.PROCESSING)
             runCurrent()
             assertThat(commandInfo.status).isEqualTo(CommandStatus.COMPLETED)
         }
@@ -114,7 +125,7 @@
             val commandInfo: CommandInfo = sut.addCommand(commandType)!!
             assertThat(commandInfo.status).isEqualTo(CommandStatus.IDLE)
 
-            runCurrent()
+            executePendingCommands()
             assertThat(commandInfo.status).isEqualTo(CommandStatus.PROCESSING)
 
             advanceTimeBy(200L)
@@ -135,12 +146,14 @@
             val commandInfo2: CommandInfo = sut.addCommand(commandType2)!!
             assertThat(commandInfo2.status).isEqualTo(CommandStatus.IDLE)
 
-            runCurrent()
+            executePendingCommands()
             assertThat(commandInfo1.status).isEqualTo(CommandStatus.PROCESSING)
             assertThat(commandInfo2.status).isEqualTo(CommandStatus.IDLE)
 
             advanceTimeBy(101L)
             assertThat(commandInfo1.status).isEqualTo(CommandStatus.COMPLETED)
+
+            executePendingCommands()
             assertThat(commandInfo2.status).isEqualTo(CommandStatus.PROCESSING)
 
             advanceTimeBy(101L)
@@ -161,12 +174,14 @@
             val commandInfo2: CommandInfo = sut.addCommand(commandType2)!!
             assertThat(commandInfo2.status).isEqualTo(CommandStatus.IDLE)
 
-            runCurrent()
+            executePendingCommands()
             assertThat(commandInfo1.status).isEqualTo(CommandStatus.PROCESSING)
             assertThat(commandInfo2.status).isEqualTo(CommandStatus.IDLE)
 
             advanceTimeBy(QUEUE_TIMEOUT)
             assertThat(commandInfo1.status).isEqualTo(CommandStatus.CANCELED)
+
+            executePendingCommands()
             assertThat(commandInfo2.status).isEqualTo(CommandStatus.PROCESSING)
 
             advanceTimeBy(101)
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/inputconsumers/NavHandleLongPressHandlerTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/inputconsumers/NavHandleLongPressHandlerTest.java
new file mode 100644
index 0000000..9018775
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/inputconsumers/NavHandleLongPressHandlerTest.java
@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.inputconsumers;
+
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+import android.content.Context;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import com.android.quickstep.DeviceConfigWrapper;
+import com.android.quickstep.NavHandle;
+import com.android.quickstep.util.TestExtensions;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class NavHandleLongPressHandlerTest {
+
+    private NavHandleLongPressHandler mLongPressHandler;
+    @Mock private NavHandle mNavHandle;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        Context context = InstrumentationRegistry.getInstrumentation().getTargetContext();
+        mLongPressHandler = new NavHandleLongPressHandler(context);
+    }
+
+    @Test
+    public void testStartNavBarAnimation_flagDisabled() {
+        try (AutoCloseable flag = overrideAnimateLPNHFlag(false)) {
+            mLongPressHandler.startNavBarAnimation(mNavHandle);
+            verify(mNavHandle, never())
+                    .animateNavBarLongPress(anyBoolean(), anyBoolean(), anyLong());
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Test
+    public void testStartNavBarAnimation_flagEnabled() {
+        try (AutoCloseable flag = overrideAnimateLPNHFlag(true)) {
+            mLongPressHandler.startNavBarAnimation(mNavHandle);
+            verify(mNavHandle).animateNavBarLongPress(anyBoolean(), anyBoolean(), anyLong());
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private AutoCloseable overrideAnimateLPNHFlag(boolean value) {
+        return TestExtensions.overrideNavConfigFlag(
+                "ANIMATE_LPNH", value, () -> DeviceConfigWrapper.get().getAnimateLpnh());
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
index a87465f..d55f2e3 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
@@ -22,7 +22,6 @@
 import android.graphics.drawable.Drawable
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.launcher3.util.TestDispatcherProvider
-import com.android.quickstep.task.thumbnail.TaskThumbnailViewModelTest
 import com.android.quickstep.util.DesktopTask
 import com.android.quickstep.util.GroupTask
 import com.android.systemui.shared.recents.model.Task
@@ -286,9 +285,14 @@
 
     private fun createThumbnailData(): ThumbnailData {
         val bitmap = mock<Bitmap>()
-        whenever(bitmap.width).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_WIDTH)
-        whenever(bitmap.height).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_HEIGHT)
+        whenever(bitmap.width).thenReturn(THUMBNAIL_WIDTH)
+        whenever(bitmap.height).thenReturn(THUMBNAIL_HEIGHT)
 
         return ThumbnailData(thumbnail = bitmap)
     }
+
+    companion object {
+        const val THUMBNAIL_WIDTH = 100
+        const val THUMBNAIL_HEIGHT = 200
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
index 33d96a8..829987c 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
@@ -22,7 +22,6 @@
 import android.graphics.Color
 import android.view.Surface
 import com.android.quickstep.recents.data.FakeTasksRepository
-import com.android.quickstep.task.thumbnail.TaskThumbnailViewModelTest
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
@@ -126,9 +125,14 @@
 
     private fun createThumbnailData(rotation: Int = Surface.ROTATION_0): ThumbnailData {
         val bitmap = mock<Bitmap>()
-        whenever(bitmap.width).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_WIDTH)
-        whenever(bitmap.height).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_HEIGHT)
+        whenever(bitmap.width).thenReturn(THUMBNAIL_WIDTH)
+        whenever(bitmap.height).thenReturn(THUMBNAIL_HEIGHT)
 
         return ThumbnailData(thumbnail = bitmap, rotation = rotation)
     }
+
+    companion object {
+        const val THUMBNAIL_WIDTH = 100
+        const val THUMBNAIL_HEIGHT = 200
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
similarity index 97%
rename from quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
rename to quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
index fcf4e56..c88a3fc 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
@@ -36,7 +36,7 @@
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.SnapshotSplash
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
 import com.android.quickstep.task.viewmodel.TaskContainerData
-import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
+import com.android.quickstep.task.viewmodel.TaskThumbnailViewModelImpl
 import com.android.quickstep.task.viewmodel.TaskViewData
 import com.android.quickstep.views.TaskViewType
 import com.android.systemui.shared.recents.model.Task
@@ -51,7 +51,7 @@
 
 /** Test for [TaskThumbnailView] */
 @RunWith(AndroidJUnit4::class)
-class TaskThumbnailViewModelTest {
+class TaskThumbnailViewModelImplTest {
     private var taskViewType = TaskViewType.SINGLE
     private val recentsViewData = RecentsViewData()
     private val taskViewData by lazy { TaskViewData(taskViewType) }
@@ -60,7 +60,7 @@
     private val mGetThumbnailPositionUseCase = mock<GetThumbnailPositionUseCase>()
     private val splashAlphaUseCase: SplashAlphaUseCase = mock()
     private val systemUnderTest by lazy {
-        TaskThumbnailViewModel(
+        TaskThumbnailViewModelImpl(
             recentsViewData,
             taskViewData,
             taskContainerData,
@@ -109,7 +109,7 @@
                         bitmap = expectedThumbnailData.thumbnail!!,
                         thumbnailRotation = Surface.ROTATION_0,
                     ),
-                    expectedIconData.icon
+                    expectedIconData.icon,
                 )
             )
     }
@@ -204,7 +204,7 @@
                         bitmap = expectedThumbnailData.thumbnail!!,
                         thumbnailRotation = Surface.ROTATION_270,
                     ),
-                    expectedIconData.icon
+                    expectedIconData.icon,
                 )
             )
     }
@@ -230,7 +230,7 @@
                         bitmap = expectedThumbnailData.thumbnail!!,
                         thumbnailRotation = Surface.ROTATION_0,
                     ),
-                    expectedIconData.icon
+                    expectedIconData.icon,
                 )
             )
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/ContextualSearchInvokerTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/ContextualSearchInvokerTest.java
new file mode 100644
index 0000000..543ffe6
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/ContextualSearchInvokerTest.java
@@ -0,0 +1,253 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.util;
+
+import static android.app.contextualsearch.ContextualSearchManager.FEATURE_CONTEXTUAL_SEARCH;
+
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_ASSISTANT_FAILED_SERVICE_ERROR;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_KEYGUARD;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_NOTIFICATION_SHADE;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_ATTEMPTED_SPLITSCREEN;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_FAILED_NOT_AVAILABLE;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_LAUNCH_OMNI_FAILED_SETTING_DISABLED;
+import static com.android.quickstep.util.ContextualSearchInvoker.KEYGUARD_SHOWING_SYSUI_FLAGS;
+import static com.android.quickstep.util.ContextualSearchInvoker.SHADE_EXPANDED_SYSUI_FLAGS;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+import android.app.contextualsearch.ContextualSearchManager;
+import android.content.Context;
+import android.content.pm.PackageManager;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.logging.StatsLogManager;
+import com.android.quickstep.DeviceConfigWrapper;
+import com.android.quickstep.SystemUiProxy;
+import com.android.quickstep.TopTaskTracker;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Robolectric unit tests for {@link ContextualSearchInvoker}
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class ContextualSearchInvokerTest {
+
+    private static final int CONTEXTUAL_SEARCH_ENTRY_POINT = 123;
+
+    private @Mock PackageManager mMockPackageManager;
+    private @Mock ContextualSearchStateManager mMockStateManager;
+    private @Mock TopTaskTracker mMockTopTaskTracker;
+    private @Mock SystemUiProxy mMockSystemUiProxy;
+    private @Mock StatsLogManager mMockStatsLogManager;
+    private @Mock StatsLogManager.StatsLogger mMockStatsLogger;
+    private @Mock ContextualSearchHapticManager mMockContextualSearchHapticManager;
+    private @Mock ContextualSearchManager mMockContextualSearchManager;
+    private ContextualSearchInvoker mContextualSearchInvoker;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        when(mMockPackageManager.hasSystemFeature(FEATURE_CONTEXTUAL_SEARCH)).thenReturn(true);
+        Context context = spy(getApplicationContext());
+        doReturn(mMockPackageManager).when(context).getPackageManager();
+        when(mMockSystemUiProxy.getLastSystemUiStateFlags()).thenReturn(0L);
+        when(mMockTopTaskTracker.getRunningSplitTaskIds()).thenReturn(new int[]{});
+        when(mMockStateManager.isContextualSearchIntentAvailable()).thenReturn(true);
+        when(mMockStateManager.isContextualSearchSettingEnabled()).thenReturn(true);
+        when(mMockStatsLogManager.logger()).thenReturn(mMockStatsLogger);
+
+        mContextualSearchInvoker = new ContextualSearchInvoker(context, mMockStateManager,
+                mMockTopTaskTracker, mMockSystemUiProxy, mMockStatsLogManager,
+                mMockContextualSearchHapticManager, mMockContextualSearchManager);
+    }
+
+    @Test
+    public void runContextualSearchInvocationChecksAndLogFailures_contextualSearchFeatureIsNotAvailable() {
+        when(mMockPackageManager.hasSystemFeature(FEATURE_CONTEXTUAL_SEARCH)).thenReturn(false);
+
+        assertFalse("Expected invocation to fail when feature is unavailable",
+                mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures());
+
+        verify(mMockStatsLogger).log(LAUNCHER_LAUNCH_ASSISTANT_FAILED_SERVICE_ERROR);
+    }
+
+    @Test
+    public void runContextualSearchInvocationChecksAndLogFailures_contextualSearchIntentIsAvailable() {
+        assertTrue("Expected invocation checks to succeed",
+                mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures());
+
+        verifyNoMoreInteractions(mMockStatsLogManager);
+    }
+
+    @Test
+    public void runContextualSearchInvocationChecksAndLogFailures_contextualSearchIntentIsNotAvailable() {
+        when(mMockStateManager.isContextualSearchIntentAvailable()).thenReturn(false);
+
+        assertFalse("Expected invocation to fail when feature is unavailable",
+                mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures());
+
+        verify(mMockStatsLogger).log(LAUNCHER_LAUNCH_OMNI_FAILED_NOT_AVAILABLE);
+    }
+
+    @Test
+    public void runContextualSearchInvocationChecksAndLogFailures_settingDisabled() {
+        when(mMockStateManager.isContextualSearchSettingEnabled()).thenReturn(false);
+
+        assertFalse("Expected invocation checks to fail when setting is disabled",
+                mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures());
+
+        verify(mMockStatsLogger).log(LAUNCHER_LAUNCH_OMNI_FAILED_SETTING_DISABLED);
+    }
+
+    @Test
+    public void runContextualSearchInvocationChecksAndLogFailures_notificationShadeIsShowing() {
+        when(mMockSystemUiProxy.getLastSystemUiStateFlags()).thenReturn(SHADE_EXPANDED_SYSUI_FLAGS);
+
+        assertFalse("Expected invocation checks to fail when notification shade is showing",
+                mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures());
+
+        verify(mMockStatsLogger).log(LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_NOTIFICATION_SHADE);
+    }
+
+    @Test
+    public void runContextualSearchInvocationChecksAndLogFailures_keyguardIsShowing() {
+        when(mMockSystemUiProxy.getLastSystemUiStateFlags()).thenReturn(
+                KEYGUARD_SHOWING_SYSUI_FLAGS);
+
+        assertFalse("Expected invocation checks to fail when keyguard is showing",
+                mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures());
+
+        verify(mMockStatsLogger).log(LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_KEYGUARD);
+    }
+
+    @Test
+    public void runContextualSearchInvocationChecksAndLogFailures_isInSplitScreen_disallowed() {
+        when(mMockStateManager.isInvocationAllowedInSplitscreen()).thenReturn(false);
+        when(mMockTopTaskTracker.getRunningSplitTaskIds()).thenReturn(new int[]{1, 2, 3});
+
+        assertFalse("Expected invocation checks to fail over split screen",
+                mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures());
+
+        // Attempt is logged regardless.
+        verify(mMockStatsLogger).log(LAUNCHER_LAUNCH_OMNI_ATTEMPTED_SPLITSCREEN);
+    }
+
+    @Test
+    public void runContextualSearchInvocationChecksAndLogFailures_isInSplitScreen_allowed() {
+        when(mMockStateManager.isInvocationAllowedInSplitscreen()).thenReturn(true);
+        when(mMockTopTaskTracker.getRunningSplitTaskIds()).thenReturn(new int[]{1, 2, 3});
+
+        assertTrue("Expected invocation checks to succeed over split screen",
+                mContextualSearchInvoker.runContextualSearchInvocationChecksAndLogFailures());
+
+        // Attempt is logged regardless.
+        verify(mMockStatsLogger).log(LAUNCHER_LAUNCH_OMNI_ATTEMPTED_SPLITSCREEN);
+    }
+
+    @Test
+    public void invokeContextualSearchUncheckedWithHaptic_cssIsAvailable_commitHapticEnabled() {
+        try (AutoCloseable flag = overrideSearchHapticCommitFlag(true)) {
+            assertTrue("Expected invocation unchecked to succeed",
+                    mContextualSearchInvoker.invokeContextualSearchUncheckedWithHaptic(
+                            CONTEXTUAL_SEARCH_ENTRY_POINT));
+            verify(mMockContextualSearchHapticManager).vibrateForSearch();
+            verify(mMockContextualSearchManager).startContextualSearch(
+                    CONTEXTUAL_SEARCH_ENTRY_POINT);
+            verifyNoMoreInteractions(mMockStatsLogManager);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Test
+    public void invokeContextualSearchUncheckedWithHaptic_cssIsAvailable_commitHapticDisabled() {
+        try (AutoCloseable flag = overrideSearchHapticCommitFlag(false)) {
+            assertTrue("Expected invocation unchecked to succeed",
+                    mContextualSearchInvoker.invokeContextualSearchUncheckedWithHaptic(
+                            CONTEXTUAL_SEARCH_ENTRY_POINT));
+            verify(mMockContextualSearchHapticManager, never()).vibrateForSearch();
+            verify(mMockContextualSearchManager).startContextualSearch(
+                    CONTEXTUAL_SEARCH_ENTRY_POINT);
+            verifyNoMoreInteractions(mMockStatsLogManager);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Test
+    public void invokeContextualSearchUncheckedWithHaptic_cssIsNotAvailable_commitHapticEnabled() {
+        when(mMockStateManager.isContextualSearchIntentAvailable()).thenReturn(false);
+
+        try (AutoCloseable flag = overrideSearchHapticCommitFlag(true)) {
+            // Still expect true since this method doesn't run the checks.
+            assertTrue("Expected invocation unchecked to succeed",
+                    mContextualSearchInvoker.invokeContextualSearchUncheckedWithHaptic(
+                            CONTEXTUAL_SEARCH_ENTRY_POINT));
+            // Still vibrate based on the flag.
+            verify(mMockContextualSearchHapticManager).vibrateForSearch();
+            verify(mMockContextualSearchManager).startContextualSearch(
+                    CONTEXTUAL_SEARCH_ENTRY_POINT);
+            verifyNoMoreInteractions(mMockStatsLogManager);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+
+    @Test
+    public void invokeContextualSearchUncheckedWithHaptic_cssIsNotAvailable_commitHapticDisabled() {
+        when(mMockStateManager.isContextualSearchIntentAvailable()).thenReturn(false);
+
+        try (AutoCloseable flag = overrideSearchHapticCommitFlag(false)) {
+            // Still expect true since this method doesn't run the checks.
+            assertTrue("Expected ContextualSearch invocation unchecked to succeed",
+                    mContextualSearchInvoker.invokeContextualSearchUncheckedWithHaptic(
+                            CONTEXTUAL_SEARCH_ENTRY_POINT));
+            // Still don't vibrate based on the flag.
+            verify(mMockContextualSearchHapticManager, never()).vibrateForSearch();
+            verify(mMockContextualSearchManager).startContextualSearch(
+                    CONTEXTUAL_SEARCH_ENTRY_POINT);
+            verifyNoMoreInteractions(mMockStatsLogManager);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private AutoCloseable overrideSearchHapticCommitFlag(boolean value) {
+        return TestExtensions.overrideNavConfigFlag(
+                "ENABLE_SEARCH_HAPTIC_COMMIT",
+                value,
+                () -> DeviceConfigWrapper.get().getEnableSearchHapticCommit());
+    }
+}
diff --git a/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
new file mode 100644
index 0000000..8968b9c
--- /dev/null
+++ b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
@@ -0,0 +1,193 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep
+
+import android.content.ComponentName
+import android.content.Intent
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import com.android.dx.mockito.inline.extended.ExtendedMockito
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
+import com.android.dx.mockito.inline.extended.StaticMockitoSession
+import com.android.launcher3.AbstractFloatingView
+import com.android.launcher3.AbstractFloatingViewHelper
+import com.android.launcher3.Flags.enableRefactorTaskThumbnail
+import com.android.launcher3.logging.StatsLogManager
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.uioverrides.QuickstepLauncher
+import com.android.launcher3.util.SplitConfigurationOptions
+import com.android.launcher3.util.TransformingTouchDelegate
+import com.android.quickstep.TaskOverlayFactory.TaskOverlay
+import com.android.quickstep.task.thumbnail.TaskThumbnailView
+import com.android.quickstep.views.LauncherRecentsView
+import com.android.quickstep.views.TaskContainer
+import com.android.quickstep.views.TaskThumbnailViewDeprecated
+import com.android.quickstep.views.TaskView
+import com.android.quickstep.views.TaskViewIcon
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.Task.TaskKey
+import com.android.window.flags.Flags
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+import org.mockito.quality.Strictness
+
+/** Test for ExternalDisplaySystemShortcut */
+class ExternalDisplaySystemShortcutTest {
+
+    @get:Rule val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
+
+    private val launcher: QuickstepLauncher = mock()
+    private val statsLogManager: StatsLogManager = mock()
+    private val statsLogger: StatsLogManager.StatsLogger = mock()
+    private val recentsView: LauncherRecentsView = mock()
+    private val taskView: TaskView = mock()
+    private val workspaceItemInfo: WorkspaceItemInfo = mock()
+    private val abstractFloatingViewHelper: AbstractFloatingViewHelper = mock()
+    private val iconView: TaskViewIcon = mock()
+    private val transformingTouchDelegate: TransformingTouchDelegate = mock()
+    private val factory: TaskShortcutFactory =
+        ExternalDisplaySystemShortcut.createFactory(abstractFloatingViewHelper)
+    private val overlayFactory: TaskOverlayFactory = mock()
+    private val overlay: TaskOverlay<*> = mock()
+
+    private lateinit var mockitoSession: StaticMockitoSession
+
+    @Before
+    fun setUp() {
+        mockitoSession =
+            mockitoSession()
+                .strictness(Strictness.LENIENT)
+                .spyStatic(DesktopModeStatus::class.java)
+                .startMocking()
+        ExtendedMockito.doReturn(true).`when` { DesktopModeStatus.enforceDeviceRestrictions() }
+        ExtendedMockito.doReturn(true).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
+        whenever(overlayFactory.createOverlay(any())).thenReturn(overlay)
+    }
+
+    @After
+    fun tearDown() {
+        mockitoSession.finishMocking()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
+    @EnableFlags(Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT)
+    fun createExternalDisplayTaskShortcut_desktopModeDisabled() {
+        val task = createTask()
+        val taskContainer = createTaskContainer(task)
+
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE,
+        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
+    )
+    fun createExternalDisplayTaskShortcut_desktopModeEnabled_deviceNotSupported() {
+        ExtendedMockito.doReturn(false).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
+
+        val taskContainer = createTaskContainer(createTask())
+
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE,
+        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
+    )
+    fun createExternalDisplayTaskShortcut_desktopModeEnabled_deviceNotSupported_overrideEnabled() {
+        ExtendedMockito.doReturn(false).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
+        ExtendedMockito.doReturn(false).`when` { DesktopModeStatus.enforceDeviceRestrictions() }
+
+        val taskContainer = spy(createTaskContainer(createTask()))
+        doReturn(workspaceItemInfo).whenever(taskContainer).itemInfo
+
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNotNull()
+    }
+
+    @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE,
+        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
+    )
+    fun externalDisplaySystemShortcutClicked() {
+        val task = createTask()
+        val taskContainer = spy(createTaskContainer(task))
+
+        whenever(launcher.getOverviewPanel<LauncherRecentsView>()).thenReturn(recentsView)
+        whenever(launcher.statsLogManager).thenReturn(statsLogManager)
+        whenever(statsLogManager.logger()).thenReturn(statsLogger)
+        whenever(statsLogger.withItemInfo(any())).thenReturn(statsLogger)
+        whenever(recentsView.moveTaskToExternalDisplay(any(), any())).thenAnswer {
+            val successCallback = it.getArgument<Runnable>(1)
+            successCallback.run()
+        }
+        doReturn(workspaceItemInfo).whenever(taskContainer).itemInfo
+
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).hasSize(1)
+        assertThat(shortcuts!!.first()).isInstanceOf(ExternalDisplaySystemShortcut::class.java)
+
+        val externalDisplayShortcut = shortcuts.first() as ExternalDisplaySystemShortcut
+
+        externalDisplayShortcut.onClick(taskView)
+
+        val allTypesExceptRebindSafe =
+            AbstractFloatingView.TYPE_ALL and AbstractFloatingView.TYPE_REBIND_SAFE.inv()
+        verify(abstractFloatingViewHelper).closeOpenViews(launcher, true, allTypesExceptRebindSafe)
+        verify(recentsView).moveTaskToExternalDisplay(eq(taskContainer), any())
+        verify(statsLogger).withItemInfo(workspaceItemInfo)
+        verify(statsLogger).log(LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_EXTERNAL_DISPLAY_TAP)
+    }
+
+    private fun createTask(): Task = Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000))
+
+    private fun createTaskContainer(task: Task): TaskContainer {
+        val snapshotView =
+            if (enableRefactorTaskThumbnail()) mock<TaskThumbnailView>()
+            else mock<TaskThumbnailViewDeprecated>()
+        return TaskContainer(
+            taskView,
+            task,
+            snapshotView,
+            iconView,
+            transformingTouchDelegate,
+            SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
+            digitalWellBeingToast = null,
+            showWindowsView = null,
+            overlayFactory,
+        )
+    }
+}
diff --git a/res/drawable/ic_private_profile_divider_badge.xml b/res/drawable/ic_private_profile_divider_badge.xml
new file mode 100644
index 0000000..07c740d
--- /dev/null
+++ b/res/drawable/ic_private_profile_divider_badge.xml
@@ -0,0 +1,26 @@
+<!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="20dp"
+    android:height="20dp"
+    android:viewportWidth="20"
+    android:viewportHeight="20">
+    <group>
+      <path
+          android:pathData="M5,9L15,9A1,1 0,0 1,16 10L16,10A1,1 0,0 1,15 11L5,11A1,1 0,0 1,4 10L4,10A1,1 0,0 1,5 9z"
+          android:fillColor="?attr/materialColorOnSurface"/>
+    </group>
+</vector>
diff --git a/res/drawable/ic_private_profile_letter_list_fast_scroller_badge.xml b/res/drawable/ic_private_profile_letter_list_fast_scroller_badge.xml
new file mode 100644
index 0000000..8d12598
--- /dev/null
+++ b/res/drawable/ic_private_profile_letter_list_fast_scroller_badge.xml
@@ -0,0 +1,28 @@
+<!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="12dp"
+    android:height="15dp"
+    android:viewportWidth="12"
+    android:viewportHeight="15">
+    <path
+        android:pathData="M5.952,0.911L0.645,2.902V6.942C0.645,10.292 2.907,13.417 5.952,14.18C8.997,13.417 11.26,10.292 11.26,6.942V2.902L5.952,0.911ZM7.943,9.536V10.863H6.616V11.526H5.289V8.103C4.333,7.818 3.63,6.942 3.63,5.887C3.63,4.607 4.672,3.565 5.952,3.565C7.233,3.565 8.274,4.607 8.274,5.887C8.274,6.935 7.571,7.818 6.616,8.103V9.536H7.943Z"
+        android:fillColor="#3C4043"
+        android:fillType="evenOdd"/>
+    <path
+        android:pathData="M5.952,6.882C6.502,6.882 6.947,6.436 6.947,5.887C6.947,5.337 6.502,4.892 5.952,4.892C5.403,4.892 4.957,5.337 4.957,5.887C4.957,6.436 5.403,6.882 5.952,6.882Z"
+        android:fillColor="#3C4043"/>
+</vector>
diff --git a/res/values-gl/strings.xml b/res/values-gl/strings.xml
index e6d6442..02389c4 100644
--- a/res/values-gl/strings.xml
+++ b/res/values-gl/strings.xml
@@ -120,7 +120,7 @@
     <string name="app_pair_name_format" msgid="8134106404716224054">"Emparellamento de aplicacións: <xliff:g id="APP1">%1$s</xliff:g> e <xliff:g id="APP2">%2$s</xliff:g>"</string>
     <string name="styles_wallpaper_button_text" msgid="8216961355289236794">"Estilo e fondo de pantalla"</string>
     <string name="edit_home_screen" msgid="8947858375782098427">"Editar pantalla de inicio"</string>
-    <string name="settings_button_text" msgid="8873672322605444408">"Axustes de Inicio"</string>
+    <string name="settings_button_text" msgid="8873672322605444408">"Configuración da pantalla de inicio"</string>
     <string name="msg_disabled_by_admin" msgid="6898038085516271325">"Función desactivada polo administrador"</string>
     <string name="allow_rotation_title" msgid="7222049633713050106">"Permitir xirar a pantalla de inicio"</string>
     <string name="allow_rotation_desc" msgid="8662546029078692509">"Ao xirar o teléfono"</string>
diff --git a/res/values-kn/strings.xml b/res/values-kn/strings.xml
index 91530b5..0c46eb4 100644
--- a/res/values-kn/strings.xml
+++ b/res/values-kn/strings.xml
@@ -120,7 +120,7 @@
     <string name="app_pair_name_format" msgid="8134106404716224054">"ಆ್ಯಪ್ ಜೋಡಿ: <xliff:g id="APP1">%1$s</xliff:g> ಮತ್ತು <xliff:g id="APP2">%2$s</xliff:g>"</string>
     <string name="styles_wallpaper_button_text" msgid="8216961355289236794">"ವಾಲ್‌ಪೇಪರ್ ಮತ್ತು ಶೈಲಿ"</string>
     <string name="edit_home_screen" msgid="8947858375782098427">"ಹೋಮ್ ಸ್ಕ್ರೀನ್ ಅನ್ನು ಎಡಿಟ್ ಮಾಡಿ"</string>
-    <string name="settings_button_text" msgid="8873672322605444408">"ಮುಖಪುಟ ಸೆಟ್ಟಿಂಗ್‌ಗಳು"</string>
+    <string name="settings_button_text" msgid="8873672322605444408">"ಹೋಮ್ ಸೆಟ್ಟಿಂಗ್‌ಗಳು"</string>
     <string name="msg_disabled_by_admin" msgid="6898038085516271325">"ನಿಮ್ಮ ನಿರ್ವಾಹಕರು ನಿಷ್ಕ್ರಿಯಗೊಳಿಸಿದ್ದಾರೆ"</string>
     <string name="allow_rotation_title" msgid="7222049633713050106">"ಹೋಮ್ ಸ್ಕ್ರೀನ್ ತಿರುಗುವಿಕೆಯನ್ನು ಅನುಮತಿಸಿ"</string>
     <string name="allow_rotation_desc" msgid="8662546029078692509">"ಫೋನ್‌ ತಿರುಗಿಸಿದಾಗ"</string>
diff --git a/src/com/android/launcher3/BaseActivity.java b/src/com/android/launcher3/BaseActivity.java
index 2e75261..3774ae3 100644
--- a/src/com/android/launcher3/BaseActivity.java
+++ b/src/com/android/launcher3/BaseActivity.java
@@ -306,6 +306,10 @@
         removeActivityFlags(ACTIVITY_STATE_RESUMED | ACTIVITY_STATE_DEFERRED_RESUMED);
     }
 
+    public boolean isPaused() {
+        return !hasBeenResumed() && (mActivityFlags & ACTIVITY_STATE_DEFERRED_RESUMED) == 0;
+    }
+
     /**
      * Sets the activity to appear as resumed.
      */
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index b0ec9b0..a8840fe 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -412,6 +412,7 @@
 
     private final List<BackPressHandler> mBackPressedHandlers = new ArrayList<>();
     private boolean mIsColdStartupAfterReboot;
+    private boolean mForceConfigUpdate;
 
     private boolean mIsNaturalScrollingEnabled;
 
@@ -756,7 +757,7 @@
     protected void onHandleConfigurationChanged() {
         Trace.beginSection("Launcher#onHandleconfigurationChanged");
         try {
-            if (!initDeviceProfile(mDeviceProfile.inv)) {
+            if (!initDeviceProfile(mDeviceProfile.inv) && !mForceConfigUpdate) {
                 return;
             }
 
@@ -770,6 +771,7 @@
             mModel.rebindCallbacks();
             updateDisallowBack();
         } finally {
+            mForceConfigUpdate = false;
             Trace.endSection();
         }
     }
@@ -1872,13 +1874,18 @@
             }
         }
 
-        // Exit spring loaded mode if necessary after adding the widget
-        Runnable onComplete = MULTI_SELECT_EDIT_MODE.get() ? null
-                : () -> mStateManager.goToState(NORMAL, SPRING_LOADED_EXIT_DELAY);
+        // Exit spring loaded mode if necessary after adding the widget; unless config activity was
+        // started.
+        Runnable onComplete = MULTI_SELECT_EDIT_MODE.get() ? null : () -> mStateManager.goToState(
+                NORMAL, SPRING_LOADED_EXIT_DELAY);
         completeAddAppWidget(appWidgetId, info, boundWidget,
                 addFlowHandler.getProviderInfo(this), addFlowHandler.needsConfigure(),
                 false, widgetPreviewBitmap);
-        mWorkspace.removeExtraEmptyScreenDelayed(delay, false, onComplete);
+        // Remove extra screen if widget drop concluded. If a config activity was started, extra
+        // screen will be removed when we get back its result.
+        if (!isActivityStarted) {
+            mWorkspace.removeExtraEmptyScreenDelayed(delay, false, onComplete);
+        }
     }
 
     public void addPendingItem(PendingAddItemInfo info, int container, int screenId,
@@ -3146,6 +3153,13 @@
         return mAnimationCoordinator;
     }
 
+    /**
+     * Set to force config update when set to true next time onHandleConfigurationChanged is called.
+     */
+    public void setForceConfigUpdate(boolean forceConfigUpdate) {
+        mForceConfigUpdate = forceConfigUpdate;
+    }
+
     @Override
     public View.OnLongClickListener getAllAppsItemLongClickListener() {
         return ItemLongClickListener.INSTANCE_ALL_APPS;
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 15641ab..42a28d6 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -176,7 +176,7 @@
                     () -> LauncherPrefs.get(mContext).removeListener(observer, THEMED_ICONS));
 
             InstallSessionTracker installSessionTracker =
-                    InstallSessionHelper.INSTANCE.get(context).registerInstallTracker(mModel);
+                    InstallSessionHelper.INSTANCE.get(context).registerInstallTracker(callbacks);
             mOnTerminateCallback.add(installSessionTracker::unregister);
         });
 
@@ -266,7 +266,7 @@
     }
 
     private class IconObserver
-            implements IconProvider.IconChangeListener, OnSharedPreferenceChangeListener {
+            implements IconProvider.IconChangeListener, LauncherPrefChangeListener {
 
         @Override
         public void onAppIconChanged(String packageName, UserHandle user) {
@@ -288,7 +288,7 @@
         }
 
         @Override
-        public void onSharedPreferenceChanged(SharedPreferences prefs, String key) {
+        public void onPrefChanged(String key) {
             if (Themes.KEY_THEMED_ICONS.equals(key)) {
                 mIconProvider.setIconThemeSupported(Themes.isThemedIconEnabled(mContext));
                 verifyIconChanged();
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
deleted file mode 100644
index 7ad17d9..0000000
--- a/src/com/android/launcher3/LauncherModel.java
+++ /dev/null
@@ -1,692 +0,0 @@
-/*
- * Copyright (C) 2008 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3;
-
-import static android.app.admin.DevicePolicyManager.ACTION_DEVICE_POLICY_RESOURCE_UPDATED;
-
-import static com.android.launcher3.LauncherAppState.ACTION_FORCE_ROLOAD;
-import static com.android.launcher3.LauncherPrefs.WORK_EDU_STEP;
-import static com.android.launcher3.config.FeatureFlags.IS_STUDIO_BUILD;
-import static com.android.launcher3.icons.cache.BaseIconCache.EMPTY_CLASS_NAME;
-import static com.android.launcher3.model.PackageUpdatedTask.OP_UPDATE;
-import static com.android.launcher3.pm.UserCache.ACTION_PROFILE_AVAILABLE;
-import static com.android.launcher3.pm.UserCache.ACTION_PROFILE_UNAVAILABLE;
-import static com.android.launcher3.testing.shared.TestProtocol.sDebugTracing;
-import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
-
-import android.content.ComponentName;
-import android.content.Context;
-import android.content.Intent;
-import android.content.pm.PackageInstaller;
-import android.content.pm.ShortcutInfo;
-import android.os.UserHandle;
-import android.text.TextUtils;
-import android.util.Log;
-import android.util.Pair;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-import androidx.annotation.WorkerThread;
-
-import com.android.launcher3.celllayout.CellPosMapper;
-import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.icons.IconCache;
-import com.android.launcher3.model.AddWorkspaceItemsTask;
-import com.android.launcher3.model.AllAppsList;
-import com.android.launcher3.model.BaseLauncherBinder;
-import com.android.launcher3.model.BgDataModel;
-import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.CacheDataUpdatedTask;
-import com.android.launcher3.model.ItemInstallQueue;
-import com.android.launcher3.model.LoaderTask;
-import com.android.launcher3.model.ModelDbController;
-import com.android.launcher3.model.ModelDelegate;
-import com.android.launcher3.model.ModelLauncherCallbacks;
-import com.android.launcher3.model.ModelTaskController;
-import com.android.launcher3.model.ModelWriter;
-import com.android.launcher3.model.PackageInstallStateChangedTask;
-import com.android.launcher3.model.PackageUpdatedTask;
-import com.android.launcher3.model.ReloadStringCacheTask;
-import com.android.launcher3.model.ShortcutsChangedTask;
-import com.android.launcher3.model.UserLockStateChangedTask;
-import com.android.launcher3.model.data.AppInfo;
-import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.pm.InstallSessionTracker;
-import com.android.launcher3.pm.PackageInstallInfo;
-import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.shortcuts.ShortcutRequest;
-import com.android.launcher3.util.ApplicationInfoWrapper;
-import com.android.launcher3.util.IntSet;
-import com.android.launcher3.util.ItemInfoMatcher;
-import com.android.launcher3.util.PackageManagerHelper;
-import com.android.launcher3.util.PackageUserKey;
-import com.android.launcher3.util.Preconditions;
-
-import java.io.FileDescriptor;
-import java.io.PrintWriter;
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.concurrent.CancellationException;
-import java.util.function.Consumer;
-import java.util.function.Supplier;
-
-/**
- * Maintains in-memory state of the Launcher. It is expected that there should be only one
- * LauncherModel object held in a static. Also provide APIs for updating the database state
- * for the Launcher.
- */
-public class LauncherModel implements InstallSessionTracker.Callback {
-    private static final boolean DEBUG_RECEIVER = false;
-
-    static final String TAG = "Launcher.Model";
-
-    @NonNull
-    private final LauncherAppState mApp;
-    @NonNull
-    private final PackageManagerHelper mPmHelper;
-    @NonNull
-    private final ModelDbController mModelDbController;
-    @NonNull
-    private final Object mLock = new Object();
-    @Nullable
-    private LoaderTask mLoaderTask;
-    private boolean mIsLoaderTaskRunning;
-
-    // only allow this once per reboot to reload work apps
-    private boolean mShouldReloadWorkProfile = true;
-
-    // Indicates whether the current model data is valid or not.
-    // We start off with everything not loaded. After that, we assume that
-    // our monitoring of the package manager provides all updates and we never
-    // need to do a requery. This is only ever touched from the loader thread.
-    private boolean mModelLoaded;
-    private boolean mModelDestroyed = false;
-    public boolean isModelLoaded() {
-        synchronized (mLock) {
-            return mModelLoaded && mLoaderTask == null && !mModelDestroyed;
-        }
-    }
-
-    @NonNull
-    private final ArrayList<Callbacks> mCallbacksList = new ArrayList<>(1);
-
-    // < only access in worker thread >
-    @NonNull
-    private final AllAppsList mBgAllAppsList;
-
-    /**
-     * All the static data should be accessed on the background thread, A lock should be acquired
-     * on this object when accessing any data from this model.
-     */
-    @NonNull
-    private final BgDataModel mBgDataModel = new BgDataModel();
-
-    @NonNull
-    private final ModelDelegate mModelDelegate;
-
-    private int mLastLoadId = -1;
-
-    // Runnable to check if the shortcuts permission has changed.
-    @NonNull
-    private final Runnable mDataValidationCheck = new Runnable() {
-        @Override
-        public void run() {
-            if (mModelLoaded) {
-                mModelDelegate.validateData();
-            }
-        }
-    };
-
-    LauncherModel(@NonNull final Context context, @NonNull final LauncherAppState app,
-            @NonNull final IconCache iconCache, @NonNull final AppFilter appFilter,
-            @NonNull final PackageManagerHelper pmHelper, final boolean isPrimaryInstance) {
-        mApp = app;
-        mPmHelper = pmHelper;
-        mModelDbController = new ModelDbController(context);
-        mBgAllAppsList = new AllAppsList(iconCache, appFilter);
-        mModelDelegate = ModelDelegate.newInstance(context, app, mPmHelper, mBgAllAppsList,
-                mBgDataModel, isPrimaryInstance);
-    }
-
-    @NonNull
-    public ModelDelegate getModelDelegate() {
-        return mModelDelegate;
-    }
-
-    public ModelDbController getModelDbController() {
-        return mModelDbController;
-    }
-
-    public ModelLauncherCallbacks newModelCallbacks() {
-        return new ModelLauncherCallbacks(this::enqueueModelUpdateTask);
-    }
-
-    /**
-     * Adds the provided items to the workspace.
-     */
-    public void addAndBindAddedWorkspaceItems(
-            @NonNull final List<Pair<ItemInfo, Object>> itemList) {
-        for (Callbacks cb : getCallbacks()) {
-            cb.preAddApps();
-        }
-        enqueueModelUpdateTask(new AddWorkspaceItemsTask(itemList));
-    }
-
-    @NonNull
-    public ModelWriter getWriter(final boolean verifyChanges, CellPosMapper cellPosMapper,
-            @Nullable final Callbacks owner) {
-        return new ModelWriter(mApp.getContext(), this, mBgDataModel, verifyChanges, cellPosMapper,
-                owner);
-    }
-
-    /**
-     * Called when the icon for an app changes, outside of package event
-     */
-    @WorkerThread
-    public void onAppIconChanged(@NonNull final String packageName,
-            @NonNull final UserHandle user) {
-        // Update the icon for the calendar package
-        Context context = mApp.getContext();
-        enqueueModelUpdateTask(new PackageUpdatedTask(OP_UPDATE, user, packageName));
-
-        List<ShortcutInfo> pinnedShortcuts = new ShortcutRequest(context, user)
-                .forPackage(packageName).query(ShortcutRequest.PINNED);
-        if (!pinnedShortcuts.isEmpty()) {
-            enqueueModelUpdateTask(new ShortcutsChangedTask(packageName, pinnedShortcuts, user,
-                    false));
-        }
-    }
-
-    /**
-     * Called when the workspace items have drastically changed
-     */
-    public void onWorkspaceUiChanged() {
-        MODEL_EXECUTOR.execute(mModelDelegate::workspaceLoadComplete);
-    }
-
-    /**
-     * Called when the model is destroyed
-     */
-    public void destroy() {
-        mModelDestroyed = true;
-        MODEL_EXECUTOR.execute(mModelDelegate::destroy);
-    }
-
-    public void onBroadcastIntent(@NonNull final Intent intent) {
-        if (DEBUG_RECEIVER || sDebugTracing) Log.d(TAG, "onReceive intent=" + intent);
-        final String action = intent.getAction();
-        if (Intent.ACTION_LOCALE_CHANGED.equals(action)) {
-            // If we have changed locale we need to clear out the labels in all apps/workspace.
-            forceReload();
-        } else if (ACTION_DEVICE_POLICY_RESOURCE_UPDATED.equals(action)) {
-            enqueueModelUpdateTask(new ReloadStringCacheTask(mModelDelegate));
-        } else if (IS_STUDIO_BUILD && ACTION_FORCE_ROLOAD.equals(action)) {
-            for (Callbacks cb : getCallbacks()) {
-                if (cb instanceof Launcher) {
-                    ((Launcher) cb).recreate();
-                }
-            }
-        }
-    }
-
-    /**
-     * Called then there use a user event
-     * @see UserCache#addUserEventListener
-     */
-    public void onUserEvent(UserHandle user, String action) {
-        if (Intent.ACTION_MANAGED_PROFILE_AVAILABLE.equals(action)
-                && mShouldReloadWorkProfile) {
-            mShouldReloadWorkProfile = false;
-            forceReload();
-        } else if (Intent.ACTION_MANAGED_PROFILE_AVAILABLE.equals(action)
-                || Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE.equals(action)) {
-            mShouldReloadWorkProfile = false;
-            enqueueModelUpdateTask(new PackageUpdatedTask(
-                    PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user));
-        } else if (UserCache.ACTION_PROFILE_LOCKED.equals(action)
-                || UserCache.ACTION_PROFILE_UNLOCKED.equals(action)) {
-            enqueueModelUpdateTask(new UserLockStateChangedTask(
-                    user, UserCache.ACTION_PROFILE_UNLOCKED.equals(action)));
-        } else if (UserCache.ACTION_PROFILE_ADDED.equals(action)
-                || UserCache.ACTION_PROFILE_REMOVED.equals(action)) {
-            forceReload();
-        } else if (ACTION_PROFILE_AVAILABLE.equals(action)
-                || ACTION_PROFILE_UNAVAILABLE.equals(action)) {
-            /*
-             * This broadcast is only available when android.os.Flags.allowPrivateProfile() is set.
-             * For Work-profile this broadcast will be sent in addition to
-             * ACTION_MANAGED_PROFILE_AVAILABLE/UNAVAILABLE.
-             * So effectively, this if block only handles the non-work profile case.
-             */
-            enqueueModelUpdateTask(new PackageUpdatedTask(
-                    PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user));
-        }
-        if (Intent.ACTION_MANAGED_PROFILE_REMOVED.equals(action)) {
-            LauncherPrefs.get(mApp.getContext()).put(WORK_EDU_STEP, 0);
-        }
-    }
-
-    /**
-     * Reloads the workspace items from the DB and re-binds the workspace. This should generally
-     * not be called as DB updates are automatically followed by UI update
-     */
-    public void forceReload() {
-        synchronized (mLock) {
-            // Stop any existing loaders first, so they don't set mModelLoaded to true later
-            stopLoader();
-            mModelLoaded = false;
-        }
-
-        // Start the loader if launcher is already running, otherwise the loader will run,
-        // the next time launcher starts
-        if (hasCallbacks()) {
-            startLoader();
-        }
-    }
-
-    /**
-     * Rebinds all existing callbacks with already loaded model
-     */
-    public void rebindCallbacks() {
-        if (hasCallbacks()) {
-            startLoader();
-        }
-    }
-
-    /**
-     * Removes an existing callback
-     */
-    public void removeCallbacks(@NonNull final Callbacks callbacks) {
-        synchronized (mCallbacksList) {
-            Preconditions.assertUIThread();
-            if (mCallbacksList.remove(callbacks)) {
-                if (stopLoader()) {
-                    // Rebind existing callbacks
-                    startLoader();
-                }
-            }
-        }
-    }
-
-    /**
-     * Adds a callbacks to receive model updates
-     * @return true if workspace load was performed synchronously
-     */
-    public boolean addCallbacksAndLoad(@NonNull final Callbacks callbacks) {
-        synchronized (mLock) {
-            addCallbacks(callbacks);
-            return startLoader(new Callbacks[] { callbacks });
-
-        }
-    }
-
-    /**
-     * Adds a callbacks to receive model updates
-     */
-    public void addCallbacks(@NonNull final Callbacks callbacks) {
-        Preconditions.assertUIThread();
-        synchronized (mCallbacksList) {
-            mCallbacksList.add(callbacks);
-        }
-    }
-
-    /**
-     * Starts the loader. Tries to bind {@params synchronousBindPage} synchronously if possible.
-     * @return true if the page could be bound synchronously.
-     */
-    public boolean startLoader() {
-        return startLoader(new Callbacks[0]);
-    }
-
-    private boolean startLoader(@NonNull final Callbacks[] newCallbacks) {
-        // Enable queue before starting loader. It will get disabled in Launcher#finishBindingItems
-        ItemInstallQueue.INSTANCE.get(mApp.getContext())
-                .pauseModelPush(ItemInstallQueue.FLAG_LOADER_RUNNING);
-        synchronized (mLock) {
-            // If there is already one running, tell it to stop.
-            boolean wasRunning = stopLoader();
-            boolean bindDirectly = mModelLoaded && !mIsLoaderTaskRunning;
-            boolean bindAllCallbacks = wasRunning || !bindDirectly || newCallbacks.length == 0;
-            final Callbacks[] callbacksList = bindAllCallbacks ? getCallbacks() : newCallbacks;
-
-            if (callbacksList.length > 0) {
-                // Clear any pending bind-runnables from the synchronized load process.
-                for (Callbacks cb : callbacksList) {
-                    MAIN_EXECUTOR.execute(cb::clearPendingBinds);
-                }
-
-                BaseLauncherBinder launcherBinder = new BaseLauncherBinder(
-                        mApp, mBgDataModel, mBgAllAppsList, callbacksList);
-                if (bindDirectly) {
-                    // Divide the set of loaded items into those that we are binding synchronously,
-                    // and everything else that is to be bound normally (asynchronously).
-                    launcherBinder.bindWorkspace(bindAllCallbacks, /* isBindSync= */ true);
-                    // For now, continue posting the binding of AllApps as there are other
-                    // issues that arise from that.
-                    launcherBinder.bindAllApps();
-                    launcherBinder.bindDeepShortcuts();
-                    launcherBinder.bindWidgets();
-                    if (FeatureFlags.CHANGE_MODEL_DELEGATE_LOADING_ORDER.get()) {
-                        mModelDelegate.bindAllModelExtras(callbacksList);
-                    }
-                    return true;
-                } else {
-                    stopLoader();
-                    mLoaderTask = new LoaderTask(
-                            mApp, mBgAllAppsList, mBgDataModel, mModelDelegate, launcherBinder);
-
-                    // Always post the loader task, instead of running directly
-                    // (even on same thread) so that we exit any nested synchronized blocks
-                    MODEL_EXECUTOR.post(mLoaderTask);
-                }
-            }
-        }
-        return false;
-    }
-
-    /**
-     * If there is already a loader task running, tell it to stop.
-     * @return true if an existing loader was stopped.
-     */
-    private boolean stopLoader() {
-        synchronized (mLock) {
-            LoaderTask oldTask = mLoaderTask;
-            mLoaderTask = null;
-            if (oldTask != null) {
-                oldTask.stopLocked();
-                return true;
-            }
-            return false;
-        }
-    }
-
-    /**
-     * Loads the model if not loaded
-     * @param callback called with the data model upon successful load or null on model thread.
-     */
-    public void loadAsync(@NonNull final Consumer<BgDataModel> callback) {
-        synchronized (mLock) {
-            if (!mModelLoaded && !mIsLoaderTaskRunning) {
-                startLoader();
-            }
-        }
-        MODEL_EXECUTOR.post(() -> callback.accept(isModelLoaded() ? mBgDataModel : null));
-    }
-
-    @Override
-    public void onInstallSessionCreated(@NonNull final PackageInstallInfo sessionInfo) {
-        if (FeatureFlags.PROMISE_APPS_IN_ALL_APPS.get()) {
-            enqueueModelUpdateTask((taskController, dataModel, apps) -> {
-                apps.addPromiseApp(mApp.getContext(), sessionInfo);
-                taskController.bindApplicationsIfNeeded();
-            });
-        }
-    }
-
-    @Override
-    public void onSessionFailure(@NonNull final String packageName,
-            @NonNull final UserHandle user) {
-        enqueueModelUpdateTask((taskController, dataModel, apps) -> {
-            IconCache iconCache = mApp.getIconCache();
-            final IntSet removedIds = new IntSet();
-            HashSet<WorkspaceItemInfo> archivedWorkspaceItemsToCacheRefresh = new HashSet<>();
-            boolean isAppArchived =
-                    new ApplicationInfoWrapper(mApp.getContext(), packageName, user).isArchived();
-            synchronized (dataModel) {
-                if (isAppArchived) {
-                    // Remove package icon cache entry for archived app in case of a session
-                    // failure.
-                    mApp.getIconCache().remove(
-                            new ComponentName(packageName, packageName + EMPTY_CLASS_NAME),
-                            user);
-                }
-
-                for (ItemInfo info : dataModel.itemsIdMap) {
-                    if (info instanceof WorkspaceItemInfo
-                            && ((WorkspaceItemInfo) info).hasPromiseIconUi()
-                            && user.equals(info.user)
-                            && info.getIntent() != null) {
-                        if (TextUtils.equals(packageName, info.getIntent().getPackage())) {
-                            removedIds.add(info.id);
-                        }
-                        if (((WorkspaceItemInfo) info).isArchived()) {
-                            WorkspaceItemInfo workspaceItem = (WorkspaceItemInfo) info;
-                            // Refresh icons on the workspace for archived apps.
-                            iconCache.getTitleAndIcon(workspaceItem,
-                                    workspaceItem.usingLowResIcon());
-                            archivedWorkspaceItemsToCacheRefresh.add(workspaceItem);
-                        }
-                    }
-                }
-
-                if (isAppArchived) {
-                    apps.updateIconsAndLabels(new HashSet<>(List.of(packageName)), user);
-                }
-            }
-
-            if (!removedIds.isEmpty() && !isAppArchived) {
-                taskController.deleteAndBindComponentsRemoved(
-                        ItemInfoMatcher.ofItemIds(removedIds),
-                        "removed because install session failed");
-            }
-            if (!archivedWorkspaceItemsToCacheRefresh.isEmpty()) {
-                taskController.bindUpdatedWorkspaceItems(
-                        archivedWorkspaceItemsToCacheRefresh.stream().toList());
-            }
-            if (isAppArchived) {
-                taskController.bindApplicationsIfNeeded();
-            }
-        });
-    }
-
-    @Override
-    public void onPackageStateChanged(@NonNull final PackageInstallInfo installInfo) {
-        enqueueModelUpdateTask(new PackageInstallStateChangedTask(installInfo));
-    }
-
-    /**
-     * Updates the icons and label of all pending icons for the provided package name.
-     */
-    @Override
-    public void onUpdateSessionDisplay(@NonNull final PackageUserKey key,
-            @NonNull final PackageInstaller.SessionInfo info) {
-        mApp.getIconCache().updateSessionCache(key, info);
-
-        HashSet<String> packages = new HashSet<>();
-        packages.add(key.mPackageName);
-        enqueueModelUpdateTask(new CacheDataUpdatedTask(
-                CacheDataUpdatedTask.OP_SESSION_UPDATE, key.mUser, packages));
-    }
-
-    public class LoaderTransaction implements AutoCloseable {
-
-        @NonNull
-        private final LoaderTask mTask;
-
-        private LoaderTransaction(@NonNull final LoaderTask task) throws CancellationException {
-            synchronized (mLock) {
-                if (mLoaderTask != task) {
-                    throw new CancellationException("Loader already stopped");
-                }
-                mLastLoadId++;
-                mTask = task;
-                mIsLoaderTaskRunning = true;
-                mModelLoaded = false;
-            }
-        }
-
-        public void commit() {
-            synchronized (mLock) {
-                // Everything loaded bind the data.
-                mModelLoaded = true;
-            }
-        }
-
-        @Override
-        public void close() {
-            synchronized (mLock) {
-                // If we are still the last one to be scheduled, remove ourselves.
-                if (mLoaderTask == mTask) {
-                    mLoaderTask = null;
-                }
-                mIsLoaderTaskRunning = false;
-            }
-        }
-    }
-
-    public LoaderTransaction beginLoader(@NonNull final LoaderTask task)
-            throws CancellationException {
-        return new LoaderTransaction(task);
-    }
-
-    /**
-     * Refreshes the cached shortcuts if the shortcut permission has changed.
-     * Current implementation simply reloads the workspace, but it can be optimized to
-     * use partial updates similar to {@link UserCache}
-     */
-    public void validateModelDataOnResume() {
-        MODEL_EXECUTOR.getHandler().removeCallbacks(mDataValidationCheck);
-        MODEL_EXECUTOR.post(mDataValidationCheck);
-    }
-
-    /**
-     * Called when the icons for packages have been updated in the icon cache.
-     */
-    public void onPackageIconsUpdated(@NonNull final HashSet<String> updatedPackages,
-            @NonNull final UserHandle user) {
-        // If any package icon has changed (app was updated while launcher was dead),
-        // update the corresponding shortcuts.
-        enqueueModelUpdateTask(new CacheDataUpdatedTask(
-                CacheDataUpdatedTask.OP_CACHE_UPDATE, user, updatedPackages));
-    }
-
-    /**
-     * Called when the labels for the widgets has updated in the icon cache.
-     */
-    public void onWidgetLabelsUpdated(@NonNull final HashSet<String> updatedPackages,
-            @NonNull final UserHandle user) {
-        enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
-            dataModel.widgetsModel.onPackageIconsUpdated(updatedPackages, user, mApp);
-            taskController.bindUpdatedWidgets(dataModel);
-        });
-    }
-
-    public void enqueueModelUpdateTask(@NonNull final ModelUpdateTask task) {
-        if (mModelDestroyed) {
-            return;
-        }
-        MODEL_EXECUTOR.execute(() -> {
-            if (!isModelLoaded()) {
-                // Loader has not yet run.
-                return;
-            }
-            ModelTaskController controller = new ModelTaskController(
-                    mApp, mBgDataModel, mBgAllAppsList, this, MAIN_EXECUTOR);
-            task.execute(controller, mBgDataModel, mBgAllAppsList);
-        });
-    }
-
-    /**
-     * A task to be executed on the current callbacks on the UI thread.
-     * If there is no current callbacks, the task is ignored.
-     */
-    public interface CallbackTask {
-
-        void execute(@NonNull Callbacks callbacks);
-    }
-
-    public interface ModelUpdateTask {
-
-        void execute(@NonNull ModelTaskController taskController,
-                @NonNull BgDataModel dataModel, @NonNull AllAppsList apps);
-    }
-
-    public void updateAndBindWorkspaceItem(@NonNull final WorkspaceItemInfo si,
-            @NonNull final ShortcutInfo info) {
-        updateAndBindWorkspaceItem(() -> {
-            si.updateFromDeepShortcutInfo(info, mApp.getContext());
-            mApp.getIconCache().getShortcutIcon(si, info);
-            return si;
-        });
-    }
-
-    /**
-     * Utility method to update a shortcut on the background thread.
-     */
-    public void updateAndBindWorkspaceItem(
-            @NonNull final Supplier<WorkspaceItemInfo> itemProvider) {
-        enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
-            WorkspaceItemInfo info = itemProvider.get();
-            taskController.getModelWriter().updateItemInDatabase(info);
-            ArrayList<WorkspaceItemInfo> update = new ArrayList<>();
-            update.add(info);
-            taskController.bindUpdatedWorkspaceItems(update);
-        });
-    }
-
-    public void refreshAndBindWidgetsAndShortcuts(@Nullable final PackageUserKey packageUser) {
-        enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
-            dataModel.widgetsModel.update(taskController.getApp(), packageUser);
-            taskController.bindUpdatedWidgets(dataModel);
-        });
-    }
-
-    public void dumpState(@Nullable final String prefix, @Nullable final FileDescriptor fd,
-            @NonNull final PrintWriter writer, @NonNull final String[] args) {
-        if (args.length > 0 && TextUtils.equals(args[0], "--all")) {
-            writer.println(prefix + "All apps list: size=" + mBgAllAppsList.data.size());
-            for (AppInfo info : mBgAllAppsList.data) {
-                writer.println(prefix + "   title=\"" + info.title
-                        + "\" bitmapIcon=" + info.bitmap.icon
-                        + " componentName=" + info.componentName.getPackageName());
-            }
-            writer.println();
-        }
-        mModelDelegate.dump(prefix, fd, writer, args);
-        mBgDataModel.dump(prefix, fd, writer, args);
-    }
-
-    /**
-     * Returns true if there are any callbacks attached to the model
-     */
-    public boolean hasCallbacks() {
-        synchronized (mCallbacksList) {
-            return !mCallbacksList.isEmpty();
-        }
-    }
-
-    /**
-     * Returns an array of currently attached callbacks
-     */
-    @NonNull
-    public Callbacks[] getCallbacks() {
-        synchronized (mCallbacksList) {
-            return mCallbacksList.toArray(new Callbacks[mCallbacksList.size()]);
-        }
-    }
-
-    /**
-     * Returns the ID for the last model load. If the load ID doesn't match for a transaction, the
-     * transaction should be ignored.
-     */
-    public int getLastLoadId() {
-        return mLastLoadId;
-    }
-}
diff --git a/src/com/android/launcher3/LauncherModel.kt b/src/com/android/launcher3/LauncherModel.kt
new file mode 100644
index 0000000..a013eaa
--- /dev/null
+++ b/src/com/android/launcher3/LauncherModel.kt
@@ -0,0 +1,499 @@
+/*
+ * Copyright (C) 2008 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3
+
+import android.app.admin.DevicePolicyManager
+import android.content.Context
+import android.content.Intent
+import android.content.pm.ShortcutInfo
+import android.os.UserHandle
+import android.text.TextUtils
+import android.util.Log
+import android.util.Pair
+import androidx.annotation.WorkerThread
+import com.android.launcher3.celllayout.CellPosMapper
+import com.android.launcher3.config.FeatureFlags
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.AddWorkspaceItemsTask
+import com.android.launcher3.model.AllAppsList
+import com.android.launcher3.model.BaseLauncherBinder
+import com.android.launcher3.model.BgDataModel
+import com.android.launcher3.model.CacheDataUpdatedTask
+import com.android.launcher3.model.ItemInstallQueue
+import com.android.launcher3.model.LoaderTask
+import com.android.launcher3.model.ModelDbController
+import com.android.launcher3.model.ModelDelegate
+import com.android.launcher3.model.ModelLauncherCallbacks
+import com.android.launcher3.model.ModelTaskController
+import com.android.launcher3.model.ModelWriter
+import com.android.launcher3.model.PackageUpdatedTask
+import com.android.launcher3.model.ReloadStringCacheTask
+import com.android.launcher3.model.ShortcutsChangedTask
+import com.android.launcher3.model.UserLockStateChangedTask
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.pm.UserCache
+import com.android.launcher3.shortcuts.ShortcutRequest
+import com.android.launcher3.testing.shared.TestProtocol.sDebugTracing
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import com.android.launcher3.util.Executors.MODEL_EXECUTOR
+import com.android.launcher3.util.PackageManagerHelper
+import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.Preconditions
+import java.io.FileDescriptor
+import java.io.PrintWriter
+import java.util.concurrent.CancellationException
+import java.util.function.Consumer
+
+/**
+ * Maintains in-memory state of the Launcher. It is expected that there should be only one
+ * LauncherModel object held in a static. Also provide APIs for updating the database state for the
+ * Launcher.
+ */
+class LauncherModel(
+    private val context: Context,
+    private val mApp: LauncherAppState,
+    private val iconCache: IconCache,
+    private val appFilter: AppFilter,
+    private val mPmHelper: PackageManagerHelper,
+    isPrimaryInstance: Boolean,
+) {
+
+    private val mCallbacksList = ArrayList<BgDataModel.Callbacks>(1)
+
+    // < only access in worker thread >
+    private val mBgAllAppsList = AllAppsList(iconCache, appFilter)
+
+    /**
+     * All the static data should be accessed on the background thread, A lock should be acquired on
+     * this object when accessing any data from this model.
+     */
+    private val mBgDataModel = BgDataModel()
+
+    val modelDelegate: ModelDelegate =
+        ModelDelegate.newInstance(
+            context,
+            mApp,
+            mPmHelper,
+            mBgAllAppsList,
+            mBgDataModel,
+            isPrimaryInstance,
+        )
+
+    val modelDbController = ModelDbController(context)
+
+    private val mLock = Any()
+
+    private var mLoaderTask: LoaderTask? = null
+    private var mIsLoaderTaskRunning = false
+
+    // only allow this once per reboot to reload work apps
+    private var mShouldReloadWorkProfile = true
+
+    // Indicates whether the current model data is valid or not.
+    // We start off with everything not loaded. After that, we assume that
+    // our monitoring of the package manager provides all updates and we never
+    // need to do a requery. This is only ever touched from the loader thread.
+    private var mModelLoaded = false
+    private var mModelDestroyed = false
+
+    fun isModelLoaded() =
+        synchronized(mLock) { mModelLoaded && mLoaderTask == null && !mModelDestroyed }
+
+    /**
+     * Returns the ID for the last model load. If the load ID doesn't match for a transaction, the
+     * transaction should be ignored.
+     */
+    var lastLoadId: Int = -1
+        private set
+
+    // Runnable to check if the shortcuts permission has changed.
+    private val mDataValidationCheck = Runnable {
+        if (mModelLoaded) {
+            modelDelegate.validateData()
+        }
+    }
+
+    fun newModelCallbacks() = ModelLauncherCallbacks(this::enqueueModelUpdateTask)
+
+    /** Adds the provided items to the workspace. */
+    fun addAndBindAddedWorkspaceItems(itemList: List<Pair<ItemInfo?, Any?>?>) {
+        callbacks.forEach { it.preAddApps() }
+        enqueueModelUpdateTask(AddWorkspaceItemsTask(itemList))
+    }
+
+    fun getWriter(
+        verifyChanges: Boolean,
+        cellPosMapper: CellPosMapper?,
+        owner: BgDataModel.Callbacks?,
+    ) = ModelWriter(mApp.context, this, mBgDataModel, verifyChanges, cellPosMapper, owner)
+
+    /** Called when the icon for an app changes, outside of package event */
+    @WorkerThread
+    fun onAppIconChanged(packageName: String, user: UserHandle) {
+        // Update the icon for the calendar package
+        enqueueModelUpdateTask(PackageUpdatedTask(PackageUpdatedTask.OP_UPDATE, user, packageName))
+        ShortcutRequest(context, user).forPackage(packageName).query(ShortcutRequest.PINNED).let {
+            if (it.isNotEmpty()) {
+                enqueueModelUpdateTask(ShortcutsChangedTask(packageName, it, user, false))
+            }
+        }
+    }
+
+    /** Called when the workspace items have drastically changed */
+    fun onWorkspaceUiChanged() {
+        MODEL_EXECUTOR.execute(modelDelegate::workspaceLoadComplete)
+    }
+
+    /** Called when the model is destroyed */
+    fun destroy() {
+        mModelDestroyed = true
+        MODEL_EXECUTOR.execute(modelDelegate::destroy)
+    }
+
+    fun onBroadcastIntent(intent: Intent) {
+        if (DEBUG_RECEIVER || sDebugTracing) Log.d(TAG, "onReceive intent=$intent")
+        when (intent.action) {
+            Intent.ACTION_LOCALE_CHANGED,
+            LauncherAppState.ACTION_FORCE_ROLOAD ->
+                // If we have changed locale we need to clear out the labels in all apps/workspace.
+                forceReload()
+            DevicePolicyManager.ACTION_DEVICE_POLICY_RESOURCE_UPDATED ->
+                enqueueModelUpdateTask(ReloadStringCacheTask(this.modelDelegate))
+        }
+    }
+
+    /**
+     * Called then there use a user event
+     *
+     * @see UserCache.addUserEventListener
+     */
+    fun onUserEvent(user: UserHandle, action: String) {
+        when (action) {
+            Intent.ACTION_MANAGED_PROFILE_AVAILABLE -> {
+                if (mShouldReloadWorkProfile) {
+                    forceReload()
+                } else {
+                    enqueueModelUpdateTask(
+                        PackageUpdatedTask(PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user)
+                    )
+                }
+                mShouldReloadWorkProfile = false
+            }
+            Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE -> {
+                mShouldReloadWorkProfile = false
+                enqueueModelUpdateTask(
+                    PackageUpdatedTask(PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user)
+                )
+            }
+            UserCache.ACTION_PROFILE_LOCKED ->
+                enqueueModelUpdateTask(UserLockStateChangedTask(user, false))
+            UserCache.ACTION_PROFILE_UNLOCKED ->
+                enqueueModelUpdateTask(UserLockStateChangedTask(user, true))
+            Intent.ACTION_MANAGED_PROFILE_REMOVED -> {
+                LauncherPrefs.get(mApp.context).put(LauncherPrefs.WORK_EDU_STEP, 0)
+                forceReload()
+            }
+            UserCache.ACTION_PROFILE_ADDED,
+            UserCache.ACTION_PROFILE_REMOVED -> forceReload()
+            UserCache.ACTION_PROFILE_AVAILABLE,
+            UserCache.ACTION_PROFILE_UNAVAILABLE -> {
+                // This broadcast is only available when android.os.Flags.allowPrivateProfile() is
+                // set. For Work-profile this broadcast will be sent in addition to
+                // ACTION_MANAGED_PROFILE_AVAILABLE/UNAVAILABLE. So effectively, this if block only
+                // handles the non-work profile case.
+                enqueueModelUpdateTask(
+                    PackageUpdatedTask(PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE, user)
+                )
+            }
+        }
+    }
+
+    /**
+     * Reloads the workspace items from the DB and re-binds the workspace. This should generally not
+     * be called as DB updates are automatically followed by UI update
+     */
+    fun forceReload() {
+        synchronized(mLock) {
+            // Stop any existing loaders first, so they don't set mModelLoaded to true later
+            stopLoader()
+            mModelLoaded = false
+        }
+        rebindCallbacks()
+    }
+
+    /** Rebinds all existing callbacks with already loaded model */
+    fun rebindCallbacks() {
+        if (hasCallbacks()) {
+            startLoader()
+        }
+    }
+
+    /** Removes an existing callback */
+    fun removeCallbacks(callbacks: BgDataModel.Callbacks) {
+        synchronized(mCallbacksList) {
+            Preconditions.assertUIThread()
+            if (mCallbacksList.remove(callbacks)) {
+                if (stopLoader()) {
+                    // Rebind existing callbacks
+                    startLoader()
+                }
+            }
+        }
+    }
+
+    /**
+     * Adds a callbacks to receive model updates
+     *
+     * @return true if workspace load was performed synchronously
+     */
+    fun addCallbacksAndLoad(callbacks: BgDataModel.Callbacks): Boolean {
+        synchronized(mLock) {
+            addCallbacks(callbacks)
+            return startLoader(arrayOf(callbacks))
+        }
+    }
+
+    /** Adds a callbacks to receive model updates */
+    fun addCallbacks(callbacks: BgDataModel.Callbacks) {
+        Preconditions.assertUIThread()
+        synchronized(mCallbacksList) { mCallbacksList.add(callbacks) }
+    }
+
+    /**
+     * Starts the loader. Tries to bind {@params synchronousBindPage} synchronously if possible.
+     *
+     * @return true if the page could be bound synchronously.
+     */
+    fun startLoader() = startLoader(arrayOf())
+
+    private fun startLoader(newCallbacks: Array<BgDataModel.Callbacks>): Boolean {
+        // Enable queue before starting loader. It will get disabled in Launcher#finishBindingItems
+        ItemInstallQueue.INSTANCE.get(context).pauseModelPush(ItemInstallQueue.FLAG_LOADER_RUNNING)
+        synchronized(mLock) {
+            // If there is already one running, tell it to stop.
+            val wasRunning = stopLoader()
+            val bindDirectly = mModelLoaded && !mIsLoaderTaskRunning
+            val bindAllCallbacks = wasRunning || !bindDirectly || newCallbacks.isEmpty()
+            val callbacksList = if (bindAllCallbacks) callbacks else newCallbacks
+            if (callbacksList.isNotEmpty()) {
+                // Clear any pending bind-runnables from the synchronized load process.
+                callbacksList.forEach { MAIN_EXECUTOR.execute(it::clearPendingBinds) }
+
+                val launcherBinder =
+                    BaseLauncherBinder(mApp, mBgDataModel, mBgAllAppsList, callbacksList)
+                if (bindDirectly) {
+                    // Divide the set of loaded items into those that we are binding synchronously,
+                    // and everything else that is to be bound normally (asynchronously).
+                    launcherBinder.bindWorkspace(bindAllCallbacks, /* isBindSync= */ true)
+                    // For now, continue posting the binding of AllApps as there are other
+                    // issues that arise from that.
+                    launcherBinder.bindAllApps()
+                    launcherBinder.bindDeepShortcuts()
+                    launcherBinder.bindWidgets()
+                    if (FeatureFlags.CHANGE_MODEL_DELEGATE_LOADING_ORDER.get()) {
+                        this.modelDelegate.bindAllModelExtras(callbacksList)
+                    }
+                    return true
+                } else {
+                    mLoaderTask =
+                        LoaderTask(
+                            mApp,
+                            mBgAllAppsList,
+                            mBgDataModel,
+                            this.modelDelegate,
+                            launcherBinder,
+                        )
+
+                    // Always post the loader task, instead of running directly
+                    // (even on same thread) so that we exit any nested synchronized blocks
+                    MODEL_EXECUTOR.post(mLoaderTask)
+                }
+            }
+        }
+        return false
+    }
+
+    /**
+     * If there is already a loader task running, tell it to stop.
+     *
+     * @return true if an existing loader was stopped.
+     */
+    private fun stopLoader(): Boolean {
+        synchronized(mLock) {
+            val oldTask: LoaderTask? = mLoaderTask
+            mLoaderTask = null
+            if (oldTask != null) {
+                oldTask.stopLocked()
+                return true
+            }
+            return false
+        }
+    }
+
+    /**
+     * Loads the model if not loaded
+     *
+     * @param callback called with the data model upon successful load or null on model thread.
+     */
+    fun loadAsync(callback: Consumer<BgDataModel?>) {
+        synchronized(mLock) {
+            if (!mModelLoaded && !mIsLoaderTaskRunning) {
+                startLoader()
+            }
+        }
+        MODEL_EXECUTOR.post { callback.accept(if (isModelLoaded()) mBgDataModel else null) }
+    }
+
+    inner class LoaderTransaction(task: LoaderTask) : AutoCloseable {
+        private var mTask: LoaderTask? = null
+
+        init {
+            synchronized(mLock) {
+                if (mLoaderTask !== task) {
+                    throw CancellationException("Loader already stopped")
+                }
+                this@LauncherModel.lastLoadId++
+                mTask = task
+                mIsLoaderTaskRunning = true
+                mModelLoaded = false
+            }
+        }
+
+        fun commit() {
+            synchronized(mLock) {
+                // Everything loaded bind the data.
+                mModelLoaded = true
+            }
+        }
+
+        override fun close() {
+            synchronized(mLock) {
+                // If we are still the last one to be scheduled, remove ourselves.
+                if (mLoaderTask === mTask) {
+                    mLoaderTask = null
+                }
+                mIsLoaderTaskRunning = false
+            }
+        }
+    }
+
+    @Throws(CancellationException::class)
+    fun beginLoader(task: LoaderTask) = LoaderTransaction(task)
+
+    /**
+     * Refreshes the cached shortcuts if the shortcut permission has changed. Current implementation
+     * simply reloads the workspace, but it can be optimized to use partial updates similar to
+     * [UserCache]
+     */
+    fun validateModelDataOnResume() {
+        MODEL_EXECUTOR.handler.removeCallbacks(mDataValidationCheck)
+        MODEL_EXECUTOR.post(mDataValidationCheck)
+    }
+
+    /** Called when the icons for packages have been updated in the icon cache. */
+    fun onPackageIconsUpdated(updatedPackages: HashSet<String?>, user: UserHandle) {
+        // If any package icon has changed (app was updated while launcher was dead),
+        // update the corresponding shortcuts.
+        enqueueModelUpdateTask(
+            CacheDataUpdatedTask(CacheDataUpdatedTask.OP_CACHE_UPDATE, user, updatedPackages)
+        )
+    }
+
+    /** Called when the labels for the widgets has updated in the icon cache. */
+    fun onWidgetLabelsUpdated(updatedPackages: HashSet<String?>, user: UserHandle) {
+        enqueueModelUpdateTask { taskController, dataModel, _ ->
+            dataModel.widgetsModel.onPackageIconsUpdated(updatedPackages, user, mApp)
+            taskController.bindUpdatedWidgets(dataModel)
+        }
+    }
+
+    fun enqueueModelUpdateTask(task: ModelUpdateTask) {
+        if (mModelDestroyed) {
+            return
+        }
+        MODEL_EXECUTOR.execute {
+            if (!isModelLoaded()) {
+                // Loader has not yet run.
+                return@execute
+            }
+            task.execute(
+                ModelTaskController(mApp, mBgDataModel, mBgAllAppsList, this, MAIN_EXECUTOR),
+                mBgDataModel,
+                mBgAllAppsList,
+            )
+        }
+    }
+
+    /**
+     * A task to be executed on the current callbacks on the UI thread. If there is no current
+     * callbacks, the task is ignored.
+     */
+    fun interface CallbackTask {
+        fun execute(callbacks: BgDataModel.Callbacks)
+    }
+
+    fun interface ModelUpdateTask {
+        fun execute(taskController: ModelTaskController, dataModel: BgDataModel, apps: AllAppsList)
+    }
+
+    fun updateAndBindWorkspaceItem(si: WorkspaceItemInfo, info: ShortcutInfo) {
+        enqueueModelUpdateTask { taskController, _, _ ->
+            si.updateFromDeepShortcutInfo(info, context)
+            iconCache.getShortcutIcon(si, info)
+            taskController.getModelWriter().updateItemInDatabase(si)
+            taskController.bindUpdatedWorkspaceItems(listOf(si))
+        }
+    }
+
+    fun refreshAndBindWidgetsAndShortcuts(packageUser: PackageUserKey?) {
+        enqueueModelUpdateTask { taskController, dataModel, _ ->
+            dataModel.widgetsModel.update(taskController.app, packageUser)
+            taskController.bindUpdatedWidgets(dataModel)
+        }
+    }
+
+    fun dumpState(prefix: String?, fd: FileDescriptor?, writer: PrintWriter, args: Array<String?>) {
+        if (args.isNotEmpty() && TextUtils.equals(args[0], "--all")) {
+            writer.println(prefix + "All apps list: size=" + mBgAllAppsList.data.size)
+            for (info in mBgAllAppsList.data) {
+                writer.println(
+                    "$prefix   title=\"${info.title}\" bitmapIcon=${info.bitmap.icon} componentName=${info.targetPackage}"
+                )
+            }
+            writer.println()
+        }
+        modelDelegate.dump(prefix, fd, writer, args)
+        mBgDataModel.dump(prefix, fd, writer, args)
+    }
+
+    /** Returns true if there are any callbacks attached to the model */
+    fun hasCallbacks() = synchronized(mCallbacksList) { mCallbacksList.isNotEmpty() }
+
+    /** Returns an array of currently attached callbacks */
+    val callbacks: Array<BgDataModel.Callbacks>
+        get() {
+            synchronized(mCallbacksList) {
+                return mCallbacksList.toTypedArray<BgDataModel.Callbacks>()
+            }
+        }
+
+    companion object {
+        private const val DEBUG_RECEIVER = false
+
+        const val TAG = "Launcher.Model"
+    }
+}
diff --git a/src/com/android/launcher3/LauncherPrefChangeListener.java b/src/com/android/launcher3/LauncherPrefChangeListener.java
new file mode 100644
index 0000000..3e9a846
--- /dev/null
+++ b/src/com/android/launcher3/LauncherPrefChangeListener.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3;
+
+import android.content.SharedPreferences;
+import android.content.SharedPreferences.OnSharedPreferenceChangeListener;
+
+/**
+ * Listener for changes in [LauncherPrefs].
+ * <p>
+ * The listener also serves as an [OnSharedPreferenceChangeListener] where
+ * [onSharedPreferenceChanged] delegates to [onPrefChanged]. Overriding [onSharedPreferenceChanged]
+ * breaks compatibility with [SharedPreferences].
+ */
+public interface LauncherPrefChangeListener extends OnSharedPreferenceChangeListener {
+
+    /** Callback invoked when the preference for [key] has changed. */
+    void onPrefChanged(String key);
+
+    @Override
+    default void onSharedPreferenceChanged(SharedPreferences sharedPreferences, String key) {
+        onPrefChanged(key);
+    }
+}
diff --git a/src/com/android/launcher3/LauncherPrefs.kt b/src/com/android/launcher3/LauncherPrefs.kt
index 13181e8..7ebfc18 100644
--- a/src/com/android/launcher3/LauncherPrefs.kt
+++ b/src/com/android/launcher3/LauncherPrefs.kt
@@ -18,7 +18,6 @@
 import android.content.Context
 import android.content.Context.MODE_PRIVATE
 import android.content.SharedPreferences
-import android.content.SharedPreferences.OnSharedPreferenceChangeListener
 import androidx.annotation.VisibleForTesting
 import com.android.launcher3.BuildConfig.WIDGET_ON_FIRST_SCREEN
 import com.android.launcher3.LauncherFiles.DEVICE_PREFERENCES_KEY
@@ -34,11 +33,177 @@
 import com.android.launcher3.util.Themes
 
 /**
- * Use same context for shared preferences, so that we use a single cached instance
+ * Manages Launcher [SharedPreferences] through [Item] instances.
  *
  * TODO(b/262721340): Replace all direct SharedPreference refs with LauncherPrefs / Item methods.
  */
-class LauncherPrefs(private val encryptedContext: Context) : SafeCloseable {
+abstract class LauncherPrefs : SafeCloseable {
+
+    /** Returns the value with type [T] for [item]. */
+    abstract fun <T> get(item: ContextualItem<T>): T
+
+    /** Returns the value with type [T] for [item]. */
+    abstract fun <T> get(item: ConstantItem<T>): T
+
+    /** Stores the values for each item in preferences. */
+    abstract fun put(vararg itemsToValues: Pair<Item, Any>)
+
+    /** Stores the [value] with type [T] for [item] in preferences. */
+    abstract fun <T : Any> put(item: Item, value: T)
+
+    /** Synchronous version of [put]. */
+    abstract fun putSync(vararg itemsToValues: Pair<Item, Any>)
+
+    /** Registers [listener] for [items]. */
+    abstract fun addListener(listener: LauncherPrefChangeListener, vararg items: Item)
+
+    /** Unregisters [listener] for [items]. */
+    abstract fun removeListener(listener: LauncherPrefChangeListener, vararg items: Item)
+
+    /** Returns `true` iff all [items] have a value. */
+    abstract fun has(vararg items: Item): Boolean
+
+    /** Removes the value for each item in [items]. */
+    abstract fun remove(vararg items: Item)
+
+    /** Synchronous version of [remove]. */
+    abstract fun removeSync(vararg items: Item)
+
+    companion object {
+        @VisibleForTesting const val BOOT_AWARE_PREFS_KEY = "boot_aware_prefs"
+
+        @JvmField
+        var INSTANCE = MainThreadInitializedObject<LauncherPrefs> { LauncherPrefsImpl(it) }
+
+        @JvmStatic fun get(context: Context): LauncherPrefs = INSTANCE.get(context)
+
+        const val TASKBAR_PINNING_KEY = "TASKBAR_PINNING_KEY"
+        const val TASKBAR_PINNING_DESKTOP_MODE_KEY = "TASKBAR_PINNING_DESKTOP_MODE_KEY"
+        const val SHOULD_SHOW_SMARTSPACE_KEY = "SHOULD_SHOW_SMARTSPACE_KEY"
+        @JvmField
+        val ICON_STATE = nonRestorableItem("pref_icon_shape_path", "", EncryptionType.ENCRYPTED)
+
+        @JvmField
+        val ENABLE_TWOLINE_ALLAPPS_TOGGLE = backedUpItem("pref_enable_two_line_toggle", false)
+        @JvmField
+        val THEMED_ICONS = backedUpItem(Themes.KEY_THEMED_ICONS, false, EncryptionType.ENCRYPTED)
+        @JvmField val PROMISE_ICON_IDS = backedUpItem(InstallSessionHelper.PROMISE_ICON_IDS, "")
+        @JvmField val WORK_EDU_STEP = backedUpItem("showed_work_profile_edu", 0)
+        @JvmField
+        val WORKSPACE_SIZE =
+            backedUpItem(DeviceGridState.KEY_WORKSPACE_SIZE, "", EncryptionType.ENCRYPTED)
+        @JvmField
+        val HOTSEAT_COUNT =
+            backedUpItem(DeviceGridState.KEY_HOTSEAT_COUNT, -1, EncryptionType.ENCRYPTED)
+        @JvmField
+        val TASKBAR_PINNING =
+            backedUpItem(TASKBAR_PINNING_KEY, false, EncryptionType.DEVICE_PROTECTED)
+        @JvmField
+        val TASKBAR_PINNING_IN_DESKTOP_MODE =
+            backedUpItem(TASKBAR_PINNING_DESKTOP_MODE_KEY, true, EncryptionType.DEVICE_PROTECTED)
+
+        @JvmField
+        val DEVICE_TYPE =
+            backedUpItem(
+                DeviceGridState.KEY_DEVICE_TYPE,
+                InvariantDeviceProfile.TYPE_PHONE,
+                EncryptionType.ENCRYPTED,
+            )
+        @JvmField
+        val DB_FILE = backedUpItem(DeviceGridState.KEY_DB_FILE, "", EncryptionType.ENCRYPTED)
+        @JvmField
+        val SHOULD_SHOW_SMARTSPACE =
+            backedUpItem(
+                SHOULD_SHOW_SMARTSPACE_KEY,
+                WIDGET_ON_FIRST_SCREEN,
+                EncryptionType.DEVICE_PROTECTED,
+            )
+        @JvmField
+        val RESTORE_DEVICE =
+            backedUpItem(
+                RestoreDbTask.RESTORED_DEVICE_TYPE,
+                InvariantDeviceProfile.TYPE_PHONE,
+                EncryptionType.ENCRYPTED,
+            )
+        @JvmField
+        val IS_FIRST_LOAD_AFTER_RESTORE =
+            nonRestorableItem(FIRST_LOAD_AFTER_RESTORE_KEY, false, EncryptionType.ENCRYPTED)
+        @JvmField val APP_WIDGET_IDS = backedUpItem(RestoreDbTask.APPWIDGET_IDS, "")
+        @JvmField val OLD_APP_WIDGET_IDS = backedUpItem(RestoreDbTask.APPWIDGET_OLD_IDS, "")
+        @JvmField
+        val GRID_NAME =
+            ConstantItem(
+                "idp_grid_name",
+                isBackedUp = true,
+                defaultValue = null,
+                encryptionType = EncryptionType.ENCRYPTED,
+                type = String::class.java,
+            )
+        @JvmField
+        val ALLOW_ROTATION =
+            backedUpItem(RotationHelper.ALLOW_ROTATION_PREFERENCE_KEY, Boolean::class.java) {
+                RotationHelper.getAllowRotationDefaultValue(DisplayController.INSTANCE.get(it).info)
+            }
+
+        // Preferences for widget configurations
+        @JvmField
+        val RECONFIGURABLE_WIDGET_EDUCATION_TIP_SEEN =
+            backedUpItem("launcher.reconfigurable_widget_education_tip_seen", false)
+
+        @JvmStatic
+        fun <T> backedUpItem(
+            sharedPrefKey: String,
+            defaultValue: T,
+            encryptionType: EncryptionType = EncryptionType.ENCRYPTED,
+        ): ConstantItem<T> =
+            ConstantItem(sharedPrefKey, isBackedUp = true, defaultValue, encryptionType)
+
+        @JvmStatic
+        fun <T> backedUpItem(
+            sharedPrefKey: String,
+            type: Class<out T>,
+            encryptionType: EncryptionType = EncryptionType.ENCRYPTED,
+            defaultValueFromContext: (c: Context) -> T,
+        ): ContextualItem<T> =
+            ContextualItem(
+                sharedPrefKey,
+                isBackedUp = true,
+                defaultValueFromContext,
+                encryptionType,
+                type,
+            )
+
+        @JvmStatic
+        fun <T> nonRestorableItem(
+            sharedPrefKey: String,
+            defaultValue: T,
+            encryptionType: EncryptionType = EncryptionType.ENCRYPTED,
+        ): ConstantItem<T> =
+            ConstantItem(sharedPrefKey, isBackedUp = false, defaultValue, encryptionType)
+
+        @Deprecated("Don't use shared preferences directly. Use other LauncherPref methods.")
+        @JvmStatic
+        fun getPrefs(context: Context): SharedPreferences {
+            // Use application context for shared preferences, so we use single cached instance
+            return context.applicationContext.getSharedPreferences(
+                SHARED_PREFERENCES_KEY,
+                MODE_PRIVATE,
+            )
+        }
+
+        @Deprecated("Don't use shared preferences directly. Use other LauncherPref methods.")
+        @JvmStatic
+        fun getDevicePrefs(context: Context): SharedPreferences {
+            // Use application context for shared preferences, so we use a single cached instance
+            return context.applicationContext.getSharedPreferences(
+                DEVICE_PREFERENCES_KEY,
+                MODE_PRIVATE,
+            )
+        }
+    }
+}
+
+private class LauncherPrefsImpl(private val encryptedContext: Context) : LauncherPrefs() {
     private val deviceProtectedStorageContext =
         encryptedContext.createDeviceProtectedStorageContext()
 
@@ -54,11 +219,11 @@
         else item.encryptedPrefs
 
     /** Wrapper around `getInner` for a `ContextualItem` */
-    fun <T> get(item: ContextualItem<T>): T =
+    override fun <T> get(item: ContextualItem<T>): T =
         getInner(item, item.defaultValueFromContext(encryptedContext))
 
     /** Wrapper around `getInner` for an `Item` */
-    fun <T> get(item: ConstantItem<T>): T = getInner(item, item.defaultValue)
+    override fun <T> get(item: ConstantItem<T>): T = getInner(item, item.defaultValue)
 
     /**
      * Retrieves the value for an [Item] from [SharedPreferences]. It handles method typing via the
@@ -97,17 +262,17 @@
      * prepareToPutValue(itemsToValues) for every distinct `SharedPreferences` file present in the
      * provided item configurations.
      */
-    fun put(vararg itemsToValues: Pair<Item, Any>): Unit =
+    override fun put(vararg itemsToValues: Pair<Item, Any>): Unit =
         prepareToPutValues(itemsToValues).forEach { it.apply() }
 
     /** See referenced `put` method above. */
-    fun <T : Any> put(item: Item, value: T): Unit = put(item.to(value))
+    override fun <T : Any> put(item: Item, value: T): Unit = put(item.to(value))
 
     /**
      * Synchronously stores all the values provided according to their associated Item
      * configuration.
      */
-    fun putSync(vararg itemsToValues: Pair<Item, Any>): Unit =
+    override fun putSync(vararg itemsToValues: Pair<Item, Any>): Unit =
         prepareToPutValues(itemsToValues).forEach { it.commit() }
 
     /**
@@ -152,7 +317,7 @@
     @Suppress("UNCHECKED_CAST")
     private fun SharedPreferences.Editor.putValue(
         item: Item,
-        value: Any?
+        value: Any?,
     ): SharedPreferences.Editor =
         when (item.type) {
             String::class.java -> putString(item.sharedPrefKey, value as? String)
@@ -176,7 +341,7 @@
      * `SharedPreferences` files associated with the provided list of items. The listener will need
      * to filter update notifications so they don't activate for non-relevant updates.
      */
-    fun addListener(listener: OnSharedPreferenceChangeListener, vararg items: Item) {
+    override fun addListener(listener: LauncherPrefChangeListener, vararg items: Item) {
         items
             .map { chooseSharedPreferences(it) }
             .distinct()
@@ -187,7 +352,7 @@
      * Stops the listener from getting notified of any more updates to any of the
      * `SharedPreferences` files associated with any of the provided list of [Item].
      */
-    fun removeListener(listener: OnSharedPreferenceChangeListener, vararg items: Item) {
+    override fun removeListener(listener: LauncherPrefChangeListener, vararg items: Item) {
         // If a listener is not registered to a SharedPreference, unregistering it does nothing
         items
             .map { chooseSharedPreferences(it) }
@@ -199,7 +364,7 @@
      * Checks if all the provided [Item] have values stored in their corresponding
      * `SharedPreferences` files.
      */
-    fun has(vararg items: Item): Boolean {
+    override fun has(vararg items: Item): Boolean {
         items
             .groupBy { chooseSharedPreferences(it) }
             .forEach { (prefs, itemsSublist) ->
@@ -211,10 +376,10 @@
     /**
      * Asynchronously removes the [Item]'s value from its corresponding `SharedPreferences` file.
      */
-    fun remove(vararg items: Item) = prepareToRemove(items).forEach { it.apply() }
+    override fun remove(vararg items: Item) = prepareToRemove(items).forEach { it.apply() }
 
     /** Synchronously removes the [Item]'s value from its corresponding `SharedPreferences` file. */
-    fun removeSync(vararg items: Item) = prepareToRemove(items).forEach { it.commit() }
+    override fun removeSync(vararg items: Item) = prepareToRemove(items).forEach { it.commit() }
 
     /**
      * Removes the key value pairs stored in `SharedPreferences` for each corresponding Item. If the
@@ -244,138 +409,6 @@
     }
 
     override fun close() {}
-
-    companion object {
-        @VisibleForTesting const val BOOT_AWARE_PREFS_KEY = "boot_aware_prefs"
-
-        @JvmField var INSTANCE = MainThreadInitializedObject { LauncherPrefs(it) }
-
-        @JvmStatic fun get(context: Context): LauncherPrefs = INSTANCE.get(context)
-
-        const val TASKBAR_PINNING_KEY = "TASKBAR_PINNING_KEY"
-        const val TASKBAR_PINNING_DESKTOP_MODE_KEY = "TASKBAR_PINNING_DESKTOP_MODE_KEY"
-        const val SHOULD_SHOW_SMARTSPACE_KEY = "SHOULD_SHOW_SMARTSPACE_KEY"
-        @JvmField
-        val ICON_STATE = nonRestorableItem("pref_icon_shape_path", "", EncryptionType.ENCRYPTED)
-
-        @JvmField
-        val ENABLE_TWOLINE_ALLAPPS_TOGGLE = backedUpItem("pref_enable_two_line_toggle", false)
-        @JvmField
-        val THEMED_ICONS = backedUpItem(Themes.KEY_THEMED_ICONS, false, EncryptionType.ENCRYPTED)
-        @JvmField val PROMISE_ICON_IDS = backedUpItem(InstallSessionHelper.PROMISE_ICON_IDS, "")
-        @JvmField val WORK_EDU_STEP = backedUpItem("showed_work_profile_edu", 0)
-        @JvmField
-        val WORKSPACE_SIZE =
-            backedUpItem(DeviceGridState.KEY_WORKSPACE_SIZE, "", EncryptionType.ENCRYPTED)
-        @JvmField
-        val HOTSEAT_COUNT =
-            backedUpItem(DeviceGridState.KEY_HOTSEAT_COUNT, -1, EncryptionType.ENCRYPTED)
-        @JvmField
-        val TASKBAR_PINNING =
-            backedUpItem(TASKBAR_PINNING_KEY, false, EncryptionType.DEVICE_PROTECTED)
-        @JvmField
-        val TASKBAR_PINNING_IN_DESKTOP_MODE =
-            backedUpItem(TASKBAR_PINNING_DESKTOP_MODE_KEY, true, EncryptionType.DEVICE_PROTECTED)
-
-        @JvmField
-        val DEVICE_TYPE =
-            backedUpItem(
-                DeviceGridState.KEY_DEVICE_TYPE,
-                InvariantDeviceProfile.TYPE_PHONE,
-                EncryptionType.ENCRYPTED
-            )
-        @JvmField
-        val DB_FILE = backedUpItem(DeviceGridState.KEY_DB_FILE, "", EncryptionType.ENCRYPTED)
-        @JvmField
-        val SHOULD_SHOW_SMARTSPACE =
-            backedUpItem(
-                SHOULD_SHOW_SMARTSPACE_KEY,
-                WIDGET_ON_FIRST_SCREEN,
-                EncryptionType.DEVICE_PROTECTED
-            )
-        @JvmField
-        val RESTORE_DEVICE =
-            backedUpItem(
-                RestoreDbTask.RESTORED_DEVICE_TYPE,
-                InvariantDeviceProfile.TYPE_PHONE,
-                EncryptionType.ENCRYPTED
-            )
-        @JvmField
-        val IS_FIRST_LOAD_AFTER_RESTORE =
-            nonRestorableItem(FIRST_LOAD_AFTER_RESTORE_KEY, false, EncryptionType.ENCRYPTED)
-        @JvmField val APP_WIDGET_IDS = backedUpItem(RestoreDbTask.APPWIDGET_IDS, "")
-        @JvmField val OLD_APP_WIDGET_IDS = backedUpItem(RestoreDbTask.APPWIDGET_OLD_IDS, "")
-        @JvmField
-        val GRID_NAME =
-            ConstantItem(
-                "idp_grid_name",
-                isBackedUp = true,
-                defaultValue = null,
-                encryptionType = EncryptionType.ENCRYPTED,
-                type = String::class.java
-            )
-        @JvmField
-        val ALLOW_ROTATION =
-            backedUpItem(RotationHelper.ALLOW_ROTATION_PREFERENCE_KEY, Boolean::class.java) {
-                RotationHelper.getAllowRotationDefaultValue(DisplayController.INSTANCE.get(it).info)
-            }
-
-        // Preferences for widget configurations
-        @JvmField
-        val RECONFIGURABLE_WIDGET_EDUCATION_TIP_SEEN =
-            backedUpItem("launcher.reconfigurable_widget_education_tip_seen", false)
-
-        @JvmStatic
-        fun <T> backedUpItem(
-            sharedPrefKey: String,
-            defaultValue: T,
-            encryptionType: EncryptionType = EncryptionType.ENCRYPTED
-        ): ConstantItem<T> =
-            ConstantItem(sharedPrefKey, isBackedUp = true, defaultValue, encryptionType)
-
-        @JvmStatic
-        fun <T> backedUpItem(
-            sharedPrefKey: String,
-            type: Class<out T>,
-            encryptionType: EncryptionType = EncryptionType.ENCRYPTED,
-            defaultValueFromContext: (c: Context) -> T
-        ): ContextualItem<T> =
-            ContextualItem(
-                sharedPrefKey,
-                isBackedUp = true,
-                defaultValueFromContext,
-                encryptionType,
-                type
-            )
-
-        @JvmStatic
-        fun <T> nonRestorableItem(
-            sharedPrefKey: String,
-            defaultValue: T,
-            encryptionType: EncryptionType = EncryptionType.ENCRYPTED
-        ): ConstantItem<T> =
-            ConstantItem(sharedPrefKey, isBackedUp = false, defaultValue, encryptionType)
-
-        @Deprecated("Don't use shared preferences directly. Use other LauncherPref methods.")
-        @JvmStatic
-        fun getPrefs(context: Context): SharedPreferences {
-            // Use application context for shared preferences, so we use single cached instance
-            return context.applicationContext.getSharedPreferences(
-                SHARED_PREFERENCES_KEY,
-                MODE_PRIVATE
-            )
-        }
-
-        @Deprecated("Don't use shared preferences directly. Use other LauncherPref methods.")
-        @JvmStatic
-        fun getDevicePrefs(context: Context): SharedPreferences {
-            // Use application context for shared preferences, so we use a single cached instance
-            return context.applicationContext.getSharedPreferences(
-                DEVICE_PREFERENCES_KEY,
-                MODE_PRIVATE
-            )
-        }
-    }
 }
 
 abstract class Item {
@@ -395,7 +428,7 @@
     val defaultValue: T,
     override val encryptionType: EncryptionType,
     // The default value can be null. If so, the type needs to be explicitly stated, or else NPE
-    override val type: Class<out T> = defaultValue!!::class.java
+    override val type: Class<out T> = defaultValue!!::class.java,
 ) : Item() {
 
     fun get(c: Context): T = LauncherPrefs.get(c).get(this)
@@ -406,7 +439,7 @@
     override val isBackedUp: Boolean,
     private val defaultSupplier: (c: Context) -> T,
     override val encryptionType: EncryptionType,
-    override val type: Class<out T>
+    override val type: Class<out T>,
 ) : Item() {
     private var default: T? = null
 
diff --git a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
index 4e1e950..e705d94 100644
--- a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
@@ -346,17 +346,12 @@
                     (LetterListTextView) LayoutInflater.from(context).inflate(
                             R.layout.fast_scroller_letter_list_text_view, mLetterList, false);
             int viewId = View.generateViewId();
-            textView.setId(viewId);
+            textView.apply(sectionInfo /* FastScrollSectionInfo */, viewId /* viewId */);
             sectionInfo.setId(viewId);
-            textView.setText(sectionInfo.sectionName);
             if (i == fastScrollSections.size() - 1) {
                 // The last section info is just a duplicate so that user can scroll to the bottom.
                 textView.setVisibility(INVISIBLE);
             }
-            ConstraintLayout.LayoutParams lp = new ConstraintLayout.LayoutParams(
-                    MATCH_CONSTRAINT, WRAP_CONTENT);
-            lp.dimensionRatio = "v,1:1";
-            textView.setLayoutParams(lp);
             textViews.add(textView);
             mLetterList.addView(textView);
         }
diff --git a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
index 8e44d65..709b52a 100644
--- a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
+++ b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
@@ -106,6 +106,7 @@
     // The of ordered component names as a result of a search query
     private final ArrayList<AdapterItem> mSearchResults = new ArrayList<>();
     private final SpannableString mPrivateProfileAppScrollerBadge;
+    private final SpannableString mPrivateProfileDividerBadge;
     private BaseAllAppsAdapter<T> mAdapter;
     private AppInfoComparator mAppNameComparator;
     private int mNumAppsPerRowAllApps;
@@ -124,9 +125,14 @@
             mAllAppsStore.addUpdateListener(this);
         }
         mPrivateProfileAppScrollerBadge = new SpannableString(" ");
-        mPrivateProfileAppScrollerBadge.setSpan(new ImageSpan(context,
+        mPrivateProfileAppScrollerBadge.setSpan(new ImageSpan(context, Flags.letterFastScroller()
+                        ? R.drawable.ic_private_profile_letter_list_fast_scroller_badge :
                         R.drawable.ic_private_profile_app_scroller_badge, ImageSpan.ALIGN_CENTER),
                 0, 1, Spannable.SPAN_EXCLUSIVE_EXCLUSIVE);
+        mPrivateProfileDividerBadge = new SpannableString(" ");
+        mPrivateProfileDividerBadge.setSpan(new ImageSpan(context,
+                        R.drawable.ic_private_profile_divider_badge, ImageSpan.ALIGN_CENTER),
+                0, 1, Spannable.SPAN_EXCLUSIVE_EXCLUSIVE);
     }
 
     /** Set the number of apps per row when device profile changes. */
@@ -404,6 +410,11 @@
         // Add system apps separator.
         if (Flags.privateSpaceSysAppsSeparation()) {
             position = mPrivateProviderManager.addSystemAppsDivider(mAdapterItems);
+            if (Flags.letterFastScroller()) {
+                FastScrollSectionInfo sectionInfo =
+                        new FastScrollSectionInfo(mPrivateProfileDividerBadge, position);
+                mFastScrollerSections.add(sectionInfo);
+            }
         }
         // Add system apps.
         position = addAppsWithSections(split.get(false), position);
@@ -437,8 +448,11 @@
                 Log.d(TAG, "addAppsWithSections: adding sectionName: " + sectionName
                     + " with appInfoTitle: " + info.title);
                 lastSectionName = sectionName;
-                mFastScrollerSections.add(new FastScrollSectionInfo(hasPrivateApps ?
-                        mPrivateProfileAppScrollerBadge : sectionName, position));
+                boolean usePrivateAppScrollerBadge = !Flags.letterFastScroller() && hasPrivateApps;
+                FastScrollSectionInfo sectionInfo = new FastScrollSectionInfo(
+                        usePrivateAppScrollerBadge ?
+                                mPrivateProfileAppScrollerBadge : sectionName, position);
+                mFastScrollerSections.add(sectionInfo);
             }
             position++;
         }
diff --git a/src/com/android/launcher3/allapps/LetterListTextView.java b/src/com/android/launcher3/allapps/LetterListTextView.java
index 9326d79..433a7f2 100644
--- a/src/com/android/launcher3/allapps/LetterListTextView.java
+++ b/src/com/android/launcher3/allapps/LetterListTextView.java
@@ -16,6 +16,9 @@
 
 package com.android.launcher3.allapps;
 
+import static androidx.constraintlayout.widget.ConstraintSet.MATCH_CONSTRAINT;
+import static androidx.constraintlayout.widget.ConstraintSet.WRAP_CONTENT;
+
 import android.content.Context;
 import android.graphics.PorterDuff;
 import android.graphics.PorterDuffColorFilter;
@@ -23,6 +26,7 @@
 import android.util.AttributeSet;
 import android.widget.TextView;
 
+import androidx.constraintlayout.widget.ConstraintLayout;
 import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.R;
@@ -71,6 +75,20 @@
     }
 
     /**
+     * Applies a viewId to the letter list text view and sets the background and text based on the
+     * sectionInfo.
+     */
+    public void apply(AlphabeticalAppsList.FastScrollSectionInfo fastScrollSectionInfo,
+            int viewId) {
+        setId(viewId);
+        setText(fastScrollSectionInfo.sectionName);
+        ConstraintLayout.LayoutParams lp = new ConstraintLayout.LayoutParams(
+                MATCH_CONSTRAINT, WRAP_CONTENT);
+        lp.dimensionRatio = "v,1:1";
+        setLayoutParams(lp);
+    }
+
+    /**
      * Animates the letter list text view based on the current finger position.
      *
      * @param currentFingerY The Y position of where the finger is placed on the fastScroller in
diff --git a/src/com/android/launcher3/dragndrop/AddItemActivity.java b/src/com/android/launcher3/dragndrop/AddItemActivity.java
index a3cfe5c..25de479 100644
--- a/src/com/android/launcher3/dragndrop/AddItemActivity.java
+++ b/src/com/android/launcher3/dragndrop/AddItemActivity.java
@@ -281,7 +281,7 @@
                 new PinShortcutRequestActivityInfo(mRequest, this);
         mWidgetCell.getWidgetView().setTag(new PendingAddShortcutInfo(shortcutInfo));
         applyWidgetItemAsync(
-                () -> new WidgetItem(shortcutInfo, mApp.getIconCache(), getPackageManager()));
+                () -> new WidgetItem(shortcutInfo, mApp.getIconCache()));
         return new PackageItemInfo(mRequest.getShortcutInfo().getPackage(),
                 mRequest.getShortcutInfo().getUserHandle());
     }
diff --git a/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java b/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java
index cc5e890..a6a50d7 100644
--- a/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java
+++ b/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java
@@ -30,7 +30,6 @@
 import android.content.Intent;
 import android.content.pm.LauncherApps;
 import android.content.pm.LauncherApps.PinItemRequest;
-import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
 import android.graphics.drawable.Drawable;
 import android.os.Build;
@@ -40,7 +39,7 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
-import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.BaseIconCache;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.pm.PinRequestHelper;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
@@ -82,12 +81,12 @@
     }
 
     @Override
-    public CharSequence getLabel(PackageManager pm) {
+    public CharSequence getLabel() {
         return mInfo.getShortLabel();
     }
 
     @Override
-    public Drawable getFullResIcon(IconCache cache) {
+    public Drawable getFullResIcon(BaseIconCache cache) {
         Drawable d = mContext.getSystemService(LauncherApps.class)
                 .getShortcutIconDrawable(mInfo, LauncherAppState.getIDP(mContext).fillResIconDpi);
         if (d == null) {
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index b6fe66a..e7c4024 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -51,8 +51,8 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.icons.cache.BaseIconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.icons.cache.CachedObjectCachingLogic;
-import com.android.launcher3.icons.cache.CachingLogic;
 import com.android.launcher3.icons.cache.LauncherActivityCachingLogic;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.AppInfo;
@@ -92,9 +92,6 @@
     private final Predicate<ItemInfoWithIcon> mIsUsingFallbackOrNonDefaultIconCheck = w ->
             w.bitmap != null && (w.bitmap.isNullOrLowRes() || !isDefaultIcon(w.bitmap, w.user));
 
-    private final CachingLogic<ComponentWithLabel> mComponentWithLabelCachingLogic;
-    private final CachingLogic<LauncherActivityInfo> mLauncherActivityInfoCachingLogic;
-
     private final LauncherApps mLauncherApps;
     private final UserCache mUserManager;
     private final InstantAppResolver mInstantAppResolver;
@@ -108,8 +105,6 @@
             IconProvider iconProvider) {
         super(context, dbFileName, MODEL_EXECUTOR.getLooper(),
                 idp.fillResIconDpi, idp.iconBitmapSize, true /* inMemoryCache */, iconProvider);
-        mComponentWithLabelCachingLogic = new CachedObjectCachingLogic(context);
-        mLauncherActivityInfoCachingLogic = LauncherActivityCachingLogic.INSTANCE;
         mLauncherApps = mContext.getSystemService(LauncherApps.class);
         mUserManager = UserCache.INSTANCE.get(mContext);
         mInstantAppResolver = InstantAppResolver.newInstance(mContext);
@@ -143,7 +138,7 @@
         removeIconsForPkg(packageName, user);
         long userSerial = mUserManager.getSerialNumberForUser(user);
         for (LauncherActivityInfo app : mLauncherApps.getActivityList(packageName, user)) {
-            addIconToDBAndMemCache(app, mLauncherActivityInfoCachingLogic, userSerial);
+            addIconToDBAndMemCache(app, LauncherActivityCachingLogic.INSTANCE, userSerial);
         }
     }
 
@@ -211,7 +206,7 @@
      */
     public synchronized void updateTitleAndIcon(AppInfo application) {
         CacheEntry entry = cacheLocked(application.componentName,
-                application.user, () -> null, mLauncherActivityInfoCachingLogic,
+                application.user, () -> null, LauncherActivityCachingLogic.INSTANCE,
                 application.usingLowResIcon() ? LookupFlag.USE_LOW_RES : LookupFlag.DEFAULT);
         if (entry.bitmap != null || !isDefaultIcon(entry.bitmap, application.user)) {
             applyCacheEntry(entry, application);
@@ -326,9 +321,9 @@
     /**
      * Loads and returns the icon for the provided object without adding it to memCache
      */
-    public synchronized String getTitleNoCache(ComponentWithLabel info) {
+    public synchronized String getTitleNoCache(CachedObject info) {
         CacheEntry entry = cacheLocked(info.getComponent(), info.getUser(), () -> info,
-                mComponentWithLabelCachingLogic,
+                CachedObjectCachingLogic.INSTANCE,
                 LookupFlag.USE_LOW_RES | LookupFlag.SKIP_ADD_TO_MEM_CACHE);
         return Utilities.trim(entry.title);
     }
@@ -344,7 +339,7 @@
         if (usePkgIcon) lookupFlags |= LookupFlag.USE_PACKAGE_ICON;
         if (useLowResIcon) lookupFlags |= LookupFlag.USE_LOW_RES;
         CacheEntry entry = cacheLocked(infoInOut.getTargetComponent(), infoInOut.user,
-                activityInfoProvider, mLauncherActivityInfoCachingLogic, lookupFlags);
+                activityInfoProvider, LauncherActivityCachingLogic.INSTANCE, lookupFlags);
         applyCacheEntry(entry, infoInOut);
     }
 
@@ -445,7 +440,7 @@
                                 cn,
                                 /* user = */ sectionKey.first,
                                 () -> duplicateIconRequests.get(0).launcherActivityInfo,
-                                mLauncherActivityInfoCachingLogic,
+                                LauncherActivityCachingLogic.INSTANCE,
                                 sectionKey.second ? LookupFlag.USE_LOW_RES : LookupFlag.DEFAULT,
                                 c);
 
@@ -494,7 +489,7 @@
                     loadFallbackIcon(
                             lai,
                             entry,
-                            mLauncherActivityInfoCachingLogic,
+                            LauncherActivityCachingLogic.INSTANCE,
                             /* usePackageIcon= */ false,
                             /* usePackageTitle= */ loadFallbackTitle,
                             cn,
@@ -504,7 +499,7 @@
                     loadFallbackTitle(
                             lai,
                             entry,
-                            mLauncherActivityInfoCachingLogic,
+                            LauncherActivityCachingLogic.INSTANCE,
                             sectionKey.first);
                 }
 
diff --git a/src/com/android/launcher3/icons/Legacy.kt b/src/com/android/launcher3/icons/Legacy.kt
deleted file mode 100644
index 3bf3bb2..0000000
--- a/src/com/android/launcher3/icons/Legacy.kt
+++ /dev/null
@@ -1,31 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3.icons
-
-import com.android.launcher3.icons.cache.CachedObject
-
-/**
- * This files contains some definitions used during refactoring to avoid breaking changes.
- *
- * TODO(b/366237794) remove this file once refactoring is complete
- */
-
-/** Temporary interface to allow easier refactoring */
-interface ComponentWithLabel : CachedObject<IconCache>
-
-/** Temporary interface to allow easier refactoring */
-interface ComponentWithLabelAndIcon : ComponentWithLabel
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index fbd24d8..e5cd76a 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -221,6 +221,9 @@
         @UiEvent(doc = "User tapped on desktop icon on a task menu.")
         LAUNCHER_SYSTEM_SHORTCUT_DESKTOP_TAP(1706),
 
+        @UiEvent(doc = "Use tapped on external display icon on a task menu,")
+        LAUNCHER_SYSTEM_SHORTCUT_EXTERNAL_DISPLAY_TAP(1957),
+
         @UiEvent(doc = "User tapped on pause app system shortcut.")
         LAUNCHER_SYSTEM_SHORTCUT_PAUSE_TAP(521),
 
@@ -798,6 +801,44 @@
         @UiEvent(doc = "User long pressed on the taskbar IME switcher button")
         LAUNCHER_TASKBAR_IME_SWITCHER_BUTTON_LONGPRESS(1798),
 
+        @UiEvent(doc = "Failed to launch assistant due to Google assistant not available")
+        LAUNCHER_LAUNCH_ASSISTANT_FAILED_NOT_AVAILABLE(1465),
+
+        @UiEvent(doc = "Failed to launch assistant due to service error")
+        LAUNCHER_LAUNCH_ASSISTANT_FAILED_SERVICE_ERROR(1466),
+
+        @UiEvent(doc = "User launched assistant by long-pressing nav handle")
+        LAUNCHER_LAUNCH_ASSISTANT_SUCCESSFUL_NAV_HANDLE(1467),
+
+        @UiEvent(doc = "Failed to launch due to Contextual Search not available")
+        LAUNCHER_LAUNCH_OMNI_FAILED_NOT_AVAILABLE(1471),
+
+        @UiEvent(doc = "Failed to launch due to Contextual Search setting disabled")
+        LAUNCHER_LAUNCH_OMNI_FAILED_SETTING_DISABLED(1632),
+
+        @UiEvent(doc = "User launched Contextual Search by long-pressing home in 3-button mode")
+        LAUNCHER_LAUNCH_OMNI_SUCCESSFUL_HOME(1481),
+
+        @UiEvent(doc = "User launched Contextual Search by using accessibility System Action")
+        LAUNCHER_LAUNCH_OMNI_SUCCESSFUL_SYSTEM_ACTION(1492),
+
+        @UiEvent(doc = "User launched Contextual Search by long pressing the meta key")
+        LAUNCHER_LAUNCH_OMNI_SUCCESSFUL_META(1606),
+
+        @UiEvent(doc = "Contextual Search invocation was attempted over the notification shade")
+        LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_NOTIFICATION_SHADE(1485),
+
+        @UiEvent(doc = "The Contextual Search all entrypoints toggle value in Settings")
+        LAUNCHER_SETTINGS_OMNI_ALL_ENTRYPOINTS_TOGGLE_VALUE(1633),
+
+        @UiEvent(doc = "Contextual Search invocation was attempted over the keyguard")
+        LAUNCHER_LAUNCH_OMNI_ATTEMPTED_OVER_KEYGUARD(1501),
+
+        @UiEvent(doc = "Contextual Search invocation was attempted while splitscreen is active")
+        LAUNCHER_LAUNCH_OMNI_ATTEMPTED_SPLITSCREEN(1505),
+
+        @UiEvent(doc = "User long press nav handle and a long press runnable was created.")
+        LAUNCHER_OMNI_GET_LONG_PRESS_RUNNABLE(1545),
         // ADD MORE
         ;
 
@@ -828,6 +869,10 @@
 
         @UiEvent(doc = "The duration of asynchronous loading workspace")
         LAUNCHER_LATENCY_STARTUP_WORKSPACE_LOADER_ASYNC(1367),
+
+        @UiEvent(doc = "Time passed between Contextual Search runnable creation and execution. This"
+                + " ensures that Recent animations have finished before Contextual Search starts.")
+        LAUNCHER_LATENCY_OMNI_RUNNABLE(1546),
         ;
 
         private final int mId;
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index dff5463..09d1146 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -72,8 +72,8 @@
 import com.android.launcher3.folder.FolderNameProvider;
 import com.android.launcher3.icons.CacheableShortcutCachingLogic;
 import com.android.launcher3.icons.CacheableShortcutInfo;
-import com.android.launcher3.icons.ComponentWithLabelAndIcon;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.icons.cache.CachedObjectCachingLogic;
 import com.android.launcher3.icons.cache.IconCacheUpdateHandler;
 import com.android.launcher3.icons.cache.LauncherActivityCachingLogic;
@@ -335,8 +335,7 @@
             verifyNotStopped();
 
             // fourth step
-            List<ComponentWithLabelAndIcon> allWidgetsList =
-                    mBgDataModel.widgetsModel.update(mApp, null);
+            List<CachedObject> allWidgetsList = mBgDataModel.widgetsModel.update(mApp, null);
             logASplit("load widgets");
 
             verifyNotStopped();
@@ -364,7 +363,7 @@
             }
 
             updateHandler.updateIcons(allWidgetsList,
-                    new CachedObjectCachingLogic(mApp.getContext()),
+                    CachedObjectCachingLogic.INSTANCE,
                     mApp.getModel()::onWidgetLabelsUpdated);
             logASplit("save widgets in icon cache");
 
diff --git a/src/com/android/launcher3/model/ModelLauncherCallbacks.kt b/src/com/android/launcher3/model/ModelLauncherCallbacks.kt
index 2ee5b80..7ba2dad 100644
--- a/src/com/android/launcher3/model/ModelLauncherCallbacks.kt
+++ b/src/com/android/launcher3/model/ModelLauncherCallbacks.kt
@@ -17,10 +17,12 @@
 package com.android.launcher3.model
 
 import android.content.pm.LauncherApps
+import android.content.pm.PackageInstaller.SessionInfo
 import android.content.pm.ShortcutInfo
 import android.os.UserHandle
 import android.text.TextUtils
 import com.android.launcher3.LauncherModel.ModelUpdateTask
+import com.android.launcher3.config.FeatureFlags
 import com.android.launcher3.logging.FileLog
 import com.android.launcher3.model.PackageUpdatedTask.OP_ADD
 import com.android.launcher3.model.PackageUpdatedTask.OP_REMOVE
@@ -28,6 +30,9 @@
 import com.android.launcher3.model.PackageUpdatedTask.OP_UNAVAILABLE
 import com.android.launcher3.model.PackageUpdatedTask.OP_UNSUSPEND
 import com.android.launcher3.model.PackageUpdatedTask.OP_UPDATE
+import com.android.launcher3.pm.InstallSessionTracker
+import com.android.launcher3.pm.PackageInstallInfo
+import com.android.launcher3.util.PackageUserKey
 import java.util.function.Consumer
 
 /**
@@ -35,7 +40,7 @@
  * model tasks
  */
 class ModelLauncherCallbacks(private var taskExecutor: Consumer<ModelUpdateTask>) :
-    LauncherApps.Callback() {
+    LauncherApps.Callback(), InstallSessionTracker.Callback {
 
     override fun onPackageAdded(packageName: String, user: UserHandle) {
         FileLog.d(TAG, "onPackageAdded triggered for packageName=$packageName, user=$user")
@@ -49,7 +54,7 @@
     override fun onPackageLoadingProgressChanged(
         packageName: String,
         user: UserHandle,
-        progress: Float
+        progress: Float,
     ) {
         taskExecutor.accept(PackageIncrementalDownloadUpdatedTask(packageName, user, progress))
     }
@@ -62,7 +67,7 @@
     override fun onPackagesAvailable(
         vararg packageNames: String,
         user: UserHandle,
-        replacing: Boolean
+        replacing: Boolean,
     ) {
         taskExecutor.accept(PackageUpdatedTask(OP_UPDATE, user, *packageNames))
     }
@@ -74,7 +79,7 @@
     override fun onPackagesUnavailable(
         packageNames: Array<String>,
         user: UserHandle,
-        replacing: Boolean
+        replacing: Boolean,
     ) {
         if (!replacing) {
             taskExecutor.accept(PackageUpdatedTask(OP_UNAVAILABLE, user, *packageNames))
@@ -88,7 +93,7 @@
     override fun onShortcutsChanged(
         packageName: String,
         shortcuts: MutableList<ShortcutInfo>,
-        user: UserHandle
+        user: UserHandle,
     ) {
         taskExecutor.accept(ShortcutsChangedTask(packageName, shortcuts, user, true))
     }
@@ -98,6 +103,37 @@
         taskExecutor.accept(PackageUpdatedTask(OP_REMOVE, user, *packages.toTypedArray()))
     }
 
+    override fun onSessionFailure(packageName: String, user: UserHandle) {
+        taskExecutor.accept(SessionFailureTask(packageName, user))
+    }
+
+    override fun onPackageStateChanged(installInfo: PackageInstallInfo) {
+        taskExecutor.accept(PackageInstallStateChangedTask(installInfo))
+    }
+
+    override fun onUpdateSessionDisplay(key: PackageUserKey, info: SessionInfo) {
+        /** Updates the icons and label of all pending icons for the provided package name. */
+        taskExecutor.accept { controller, _, _ ->
+            controller.app.iconCache.updateSessionCache(key, info)
+        }
+        taskExecutor.accept(
+            CacheDataUpdatedTask(
+                CacheDataUpdatedTask.OP_SESSION_UPDATE,
+                key.mUser,
+                hashSetOf(key.mPackageName),
+            )
+        )
+    }
+
+    override fun onInstallSessionCreated(sessionInfo: PackageInstallInfo) {
+        if (FeatureFlags.PROMISE_APPS_IN_ALL_APPS.get()) {
+            taskExecutor.accept { taskController, _, apps ->
+                apps.addPromiseApp(taskController.app.context, sessionInfo)
+                taskController.bindApplicationsIfNeeded()
+            }
+        }
+    }
+
     companion object {
         private const val TAG = "LauncherAppsCallbackImpl"
     }
diff --git a/src/com/android/launcher3/model/SessionFailureTask.kt b/src/com/android/launcher3/model/SessionFailureTask.kt
new file mode 100644
index 0000000..0d006fa
--- /dev/null
+++ b/src/com/android/launcher3/model/SessionFailureTask.kt
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.model
+
+import android.content.ComponentName
+import android.os.UserHandle
+import android.text.TextUtils
+import com.android.launcher3.LauncherModel.ModelUpdateTask
+import com.android.launcher3.icons.cache.BaseIconCache
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.util.ApplicationInfoWrapper
+import com.android.launcher3.util.ItemInfoMatcher
+
+/** Model task run when there is a package install session failure */
+class SessionFailureTask(val packageName: String, val user: UserHandle) : ModelUpdateTask {
+
+    override fun execute(
+        taskController: ModelTaskController,
+        dataModel: BgDataModel,
+        apps: AllAppsList,
+    ) {
+        val iconCache = taskController.app.iconCache
+        val isAppArchived =
+            ApplicationInfoWrapper(taskController.app.context, packageName, user).isArchived()
+        synchronized(dataModel) {
+            if (isAppArchived) {
+                val updatedItems = mutableListOf<WorkspaceItemInfo>()
+                // Remove package icon cache entry for archived app in case of a session
+                // failure.
+                iconCache.remove(
+                    ComponentName(packageName, packageName + BaseIconCache.EMPTY_CLASS_NAME),
+                    user,
+                )
+                for (info in dataModel.itemsIdMap) {
+                    if (info is WorkspaceItemInfo && info.isArchived && user == info.user) {
+                        // Refresh icons on the workspace for archived apps.
+                        iconCache.getTitleAndIcon(info, info.usingLowResIcon())
+                        updatedItems.add(info)
+                    }
+                }
+
+                if (updatedItems.isNotEmpty()) {
+                    taskController.bindUpdatedWorkspaceItems(updatedItems)
+                }
+                apps.updateIconsAndLabels(hashSetOf(packageName), user)
+                taskController.bindApplicationsIfNeeded()
+            } else {
+                val removedItems =
+                    dataModel.itemsIdMap.filter { info ->
+                        (info is WorkspaceItemInfo && info.hasPromiseIconUi()) &&
+                            user == info.user &&
+                            TextUtils.equals(packageName, info.intent.getPackage())
+                    }
+                if (removedItems.isNotEmpty()) {
+                    taskController.deleteAndBindComponentsRemoved(
+                        ItemInfoMatcher.ofItems(removedItems),
+                        "removed because install session failed",
+                    )
+                }
+            }
+        }
+    }
+}
diff --git a/src/com/android/launcher3/model/WidgetItem.java b/src/com/android/launcher3/model/WidgetItem.java
index ac9f2d6..e757a68 100644
--- a/src/com/android/launcher3/model/WidgetItem.java
+++ b/src/com/android/launcher3/model/WidgetItem.java
@@ -7,7 +7,6 @@
 import android.annotation.SuppressLint;
 import android.content.Context;
 import android.content.pm.ActivityInfo;
-import android.content.pm.PackageManager;
 import android.content.res.Resources;
 import android.util.SparseArray;
 import android.widget.RemoteViews;
@@ -75,10 +74,10 @@
         this(info, idp, iconCache, context, new WidgetManagerHelper(context));
     }
 
-    public WidgetItem(ShortcutConfigActivityInfo info, IconCache iconCache, PackageManager pm) {
+    public WidgetItem(ShortcutConfigActivityInfo info, IconCache iconCache) {
         super(info.getComponent(), info.getUser());
         label = info.isPersistable() ? iconCache.getTitleNoCache(info) :
-                Utilities.trim(info.getLabel(pm));
+                Utilities.trim(info.getLabel());
         description = null;
         widgetInfo = null;
         activityInfo = info;
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index c949ce6..b450f46 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -14,7 +14,6 @@
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.os.UserHandle;
 import android.util.Log;
 import android.util.Pair;
@@ -27,8 +26,8 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.icons.ComponentWithLabelAndIcon;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
 import com.android.launcher3.util.ComponentKey;
@@ -96,20 +95,18 @@
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
      */
-    public List<ComponentWithLabelAndIcon> update(
+    public List<CachedObject> update(
             LauncherAppState app, @Nullable PackageUserKey packageUser) {
         if (!WIDGETS_ENABLED) {
-            return Collections.emptyList();
+            return new ArrayList<>();
         }
         Preconditions.assertWorkerThread();
 
         Context context = app.getContext();
         final ArrayList<WidgetItem> widgetsAndShortcuts = new ArrayList<>();
-        List<ComponentWithLabelAndIcon> updatedItems = new ArrayList<>();
+        List<CachedObject> updatedItems = new ArrayList<>();
         try {
             InvariantDeviceProfile idp = app.getInvariantDeviceProfile();
-            PackageManager pm = app.getContext().getPackageManager();
-
             // Widgets
             WidgetManagerHelper widgetManager = new WidgetManagerHelper(context);
             for (AppWidgetProviderInfo widgetInfo : widgetManager.getAllProviders(packageUser)) {
@@ -125,7 +122,7 @@
             // Shortcuts
             for (ShortcutConfigActivityInfo info :
                     queryList(context, packageUser)) {
-                widgetsAndShortcuts.add(new WidgetItem(info, app.getIconCache(), pm));
+                widgetsAndShortcuts.add(new WidgetItem(info, app.getIconCache()));
                 updatedItems.add(info);
             }
             setWidgetsAndShortcuts(widgetsAndShortcuts, app, packageUser);
@@ -190,8 +187,7 @@
                     WidgetItem item = items.get(i);
                     if (item.user.equals(user)) {
                         if (item.activityInfo != null) {
-                            items.set(i, new WidgetItem(item.activityInfo, app.getIconCache(),
-                                    app.getContext().getPackageManager()));
+                            items.set(i, new WidgetItem(item.activityInfo, app.getIconCache()));
                         } else {
                             items.set(i, new WidgetItem(item.widgetInfo,
                                     app.getInvariantDeviceProfile(), app.getIconCache(),
diff --git a/src/com/android/launcher3/pm/InstallSessionHelper.java b/src/com/android/launcher3/pm/InstallSessionHelper.java
index f36f595..afc5117 100644
--- a/src/com/android/launcher3/pm/InstallSessionHelper.java
+++ b/src/com/android/launcher3/pm/InstallSessionHelper.java
@@ -38,13 +38,10 @@
 import com.android.launcher3.model.ItemInstallQueue;
 import com.android.launcher3.util.ApplicationInfoWrapper;
 import com.android.launcher3.util.DaggerSingletonObject;
-import com.android.launcher3.util.DaggerSingletonTracker;
-import com.android.launcher3.util.ExecutorUtil;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
 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;
@@ -59,7 +56,7 @@
  */
 @SuppressWarnings("NewApi")
 @LauncherAppSingleton
-public class InstallSessionHelper implements SafeCloseable {
+public class InstallSessionHelper {
 
     @NonNull
     private static final String LOG = "InstallSessionHelper";
@@ -91,17 +88,12 @@
     private IntSet mPromiseIconIds;
 
     @Inject
-    public InstallSessionHelper(@NonNull @ApplicationContext final Context context,
-            DaggerSingletonTracker tracker) {
+    public InstallSessionHelper(@NonNull @ApplicationContext final Context context) {
         mInstaller = context.getPackageManager().getPackageInstaller();
         mAppContext = context.getApplicationContext();
         mLauncherApps = context.getSystemService(LauncherApps.class);
-        ExecutorUtil.executeSyncOnMainOrFail(() -> tracker.addCloseable(this));
     }
 
-    @Override
-    public void close() { }
-
     @WorkerThread
     @NonNull
     private IntSet getPromiseIconIds() {
diff --git a/src/com/android/launcher3/pm/ShortcutConfigActivityInfo.java b/src/com/android/launcher3/pm/ShortcutConfigActivityInfo.java
index 3064abf..409174e 100644
--- a/src/com/android/launcher3/pm/ShortcutConfigActivityInfo.java
+++ b/src/com/android/launcher3/pm/ShortcutConfigActivityInfo.java
@@ -29,7 +29,6 @@
 import android.content.pm.ApplicationInfo;
 import android.content.pm.LauncherActivityInfo;
 import android.content.pm.LauncherApps;
-import android.content.pm.PackageManager;
 import android.graphics.drawable.Drawable;
 import android.os.Build;
 import android.os.Process;
@@ -41,8 +40,8 @@
 
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
-import com.android.launcher3.icons.ComponentWithLabelAndIcon;
-import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.BaseIconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.ApplicationInfoWrapper;
 import com.android.launcher3.util.PackageUserKey;
@@ -54,7 +53,7 @@
 /**
  * Wrapper class for representing a shortcut configure activity.
  */
-public abstract class ShortcutConfigActivityInfo implements ComponentWithLabelAndIcon {
+public abstract class ShortcutConfigActivityInfo implements CachedObject {
 
     private static final String TAG = "SCActivityInfo";
 
@@ -91,7 +90,7 @@
     }
 
     @Override
-    public abstract Drawable getFullResIcon(IconCache cache);
+    public abstract Drawable getFullResIcon(BaseIconCache cache);
 
     /**
      * Return a WorkspaceItemInfo, if it can be created directly on drop, without requiring any
@@ -125,7 +124,7 @@
     }
 
     /**
-     * Returns true if various properties ({@link #getLabel(PackageManager)},
+     * Returns true if various properties ({@link #getLabel()},
      * {@link #getFullResIcon}) can be safely persisted.
      */
     public boolean isPersistable() {
@@ -144,12 +143,12 @@
         }
 
         @Override
-        public CharSequence getLabel(PackageManager pm) {
+        public CharSequence getLabel() {
             return mInfo.getLabel();
         }
 
         @Override
-        public Drawable getFullResIcon(IconCache cache) {
+        public Drawable getFullResIcon(BaseIconCache cache) {
             return cache.getFullResIcon(mInfo.getActivityInfo());
         }
 
diff --git a/src/com/android/launcher3/statemanager/StateManager.java b/src/com/android/launcher3/statemanager/StateManager.java
index 303290d..763f3ba 100644
--- a/src/com/android/launcher3/statemanager/StateManager.java
+++ b/src/com/android/launcher3/statemanager/StateManager.java
@@ -18,6 +18,7 @@
 
 import static android.animation.ValueAnimator.areAnimatorsEnabled;
 
+import static com.android.launcher3.Flags.enableStateManagerProtoLog;
 import static com.android.launcher3.anim.AnimatorPlaybackController.callListenerCommandRecursively;
 import static com.android.launcher3.states.StateAnimationConfig.HANDLE_STATE_APPLY;
 import static com.android.launcher3.states.StateAnimationConfig.SKIP_ALL_ANIMATIONS;
@@ -39,6 +40,7 @@
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.states.StateAnimationConfig.AnimationFlags;
 import com.android.launcher3.states.StateAnimationConfig.AnimationPropertyFlags;
+import com.android.launcher3.util.StateManagerProtoLogProxy;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -243,7 +245,10 @@
 
     private void goToState(
             STATE_TYPE state, boolean animated, long delay, AnimatorListener listener) {
-        if (DEBUG) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logGoToState(
+                    mState, state, getTrimmedStackTrace("StateManager.goToState"));
+        } else if (DEBUG) {
             Log.d(TAG, "goToState - fromState: " + mState + ", toState: " + state
                     + ", partial trace:\n" + getTrimmedStackTrace("StateManager.goToState"));
         }
@@ -331,7 +336,10 @@
      */
     public AnimatorSet createAtomicAnimation(
             STATE_TYPE fromState, STATE_TYPE toState, StateAnimationConfig config) {
-        if (DEBUG) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logCreateAtomicAnimation(
+                    mState, toState, getTrimmedStackTrace("StateManager.createAtomicAnimation"));
+        } else if (DEBUG) {
             Log.d(TAG, "createAtomicAnimation - fromState: " + fromState + ", toState: " + toState
                     + ", partial trace:\n" + getTrimmedStackTrace(
                             "StateManager.createAtomicAnimation"));
@@ -408,7 +416,9 @@
         mState = state;
         mStatefulContainer.onStateSetStart(mState);
 
-        if (DEBUG) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logOnStateTransitionStart(state);
+        } else if (DEBUG) {
             Log.d(TAG, "onStateTransitionStart - state: " + state);
         }
         for (int i = mListeners.size() - 1; i >= 0; i--) {
@@ -428,7 +438,9 @@
             setRestState(null);
         }
 
-        if (DEBUG) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logOnStateTransitionEnd(state);
+        } else if (DEBUG) {
             Log.d(TAG, "onStateTransitionEnd - state: " + state);
         }
         for (int i = mListeners.size() - 1; i >= 0; i--) {
@@ -468,7 +480,11 @@
      * Cancels the current animation.
      */
     public void cancelAnimation() {
-        if (DEBUG && mConfig.currentAnimation != null) {
+        if (enableStateManagerProtoLog()) {
+            StateManagerProtoLogProxy.logCancelAnimation(
+                    mConfig.currentAnimation != null,
+                    getTrimmedStackTrace("StateManager.cancelAnimation"));
+        } else if (DEBUG && mConfig.currentAnimation != null) {
             Log.d(TAG, "cancelAnimation - with ongoing animation"
                     + ", partial trace:\n" + getTrimmedStackTrace("StateManager.cancelAnimation"));
         }
diff --git a/src/com/android/launcher3/states/RotationHelper.java b/src/com/android/launcher3/states/RotationHelper.java
index fdb37f0..b3bcada 100644
--- a/src/com/android/launcher3/states/RotationHelper.java
+++ b/src/com/android/launcher3/states/RotationHelper.java
@@ -26,8 +26,6 @@
 import static com.android.launcher3.util.window.WindowManagerProxy.MIN_TABLET_WIDTH;
 
 import android.content.Context;
-import android.content.SharedPreferences;
-import android.content.SharedPreferences.OnSharedPreferenceChangeListener;
 import android.os.Handler;
 import android.os.Message;
 
@@ -36,13 +34,14 @@
 
 import com.android.launcher3.BaseActivity;
 import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.LauncherPrefChangeListener;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.util.DisplayController;
 
 /**
  * Utility class to manage launcher rotation
  */
-public class RotationHelper implements OnSharedPreferenceChangeListener,
+public class RotationHelper implements LauncherPrefChangeListener,
         DeviceProfile.OnDeviceProfileChangeListener,
         DisplayController.DisplayInfoChangeListener {
 
@@ -112,7 +111,7 @@
     }
 
     @Override
-    public void onSharedPreferenceChanged(SharedPreferences sharedPreferences, String s) {
+    public void onPrefChanged(String s) {
         if (mDestroyed || mIgnoreAutoRotateSettings) return;
         boolean wasRotationEnabled = mHomeRotationEnabled;
         mHomeRotationEnabled = LauncherPrefs.get(mActivity).get(ALLOW_ROTATION);
diff --git a/src/com/android/launcher3/util/DaggerSingletonObject.java b/src/com/android/launcher3/util/DaggerSingletonObject.java
index b8cf2ae..febe6af 100644
--- a/src/com/android/launcher3/util/DaggerSingletonObject.java
+++ b/src/com/android/launcher3/util/DaggerSingletonObject.java
@@ -29,7 +29,7 @@
  * We should delete this class at the end and use @Inject to get dagger provided singletons.
  */
 
-public class DaggerSingletonObject<T extends SafeCloseable> {
+public class DaggerSingletonObject<T> {
     private final Function<LauncherAppComponent, T> mFunction;
 
     public DaggerSingletonObject(Function<LauncherAppComponent, T> function) {
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index c59cc81..0b45118 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -35,7 +35,6 @@
 import android.content.ComponentCallbacks;
 import android.content.Context;
 import android.content.Intent;
-import android.content.SharedPreferences;
 import android.content.res.Configuration;
 import android.graphics.Point;
 import android.graphics.Rect;
@@ -51,6 +50,7 @@
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.InvariantDeviceProfile.DeviceType;
+import com.android.launcher3.LauncherPrefChangeListener;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
@@ -116,8 +116,7 @@
     private Info mInfo;
     private boolean mDestroyed = false;
 
-    private SharedPreferences.OnSharedPreferenceChangeListener
-            mTaskbarPinningPreferenceChangeListener;
+    private LauncherPrefChangeListener mTaskbarPinningPreferenceChangeListener;
 
     @VisibleForTesting
     protected DisplayController(Context context) {
@@ -142,19 +141,18 @@
     }
 
     private void attachTaskbarPinningSharedPreferenceChangeListener(Context context) {
-        mTaskbarPinningPreferenceChangeListener =
-                (sharedPreferences, key) -> {
-                    LauncherPrefs prefs = LauncherPrefs.get(mContext);
-                    boolean isTaskbarPinningChanged = TASKBAR_PINNING_KEY.equals(key)
-                            && mInfo.mIsTaskbarPinned != prefs.get(TASKBAR_PINNING);
-                    boolean isTaskbarPinningDesktopModeChanged =
-                            TASKBAR_PINNING_DESKTOP_MODE_KEY.equals(key)
-                                    && mInfo.mIsTaskbarPinnedInDesktopMode != prefs.get(
-                                    TASKBAR_PINNING_IN_DESKTOP_MODE);
-                    if (isTaskbarPinningChanged || isTaskbarPinningDesktopModeChanged) {
-                        notifyConfigChange();
-                    }
-                };
+        mTaskbarPinningPreferenceChangeListener = key -> {
+            LauncherPrefs prefs = LauncherPrefs.get(mContext);
+            boolean isTaskbarPinningChanged = TASKBAR_PINNING_KEY.equals(key)
+                    && mInfo.mIsTaskbarPinned != prefs.get(TASKBAR_PINNING);
+            boolean isTaskbarPinningDesktopModeChanged =
+                    TASKBAR_PINNING_DESKTOP_MODE_KEY.equals(key)
+                            && mInfo.mIsTaskbarPinnedInDesktopMode != prefs.get(
+                            TASKBAR_PINNING_IN_DESKTOP_MODE);
+            if (isTaskbarPinningChanged || isTaskbarPinningDesktopModeChanged) {
+                notifyConfigChange();
+            }
+        };
 
         LauncherPrefs.get(context).addListener(
                 mTaskbarPinningPreferenceChangeListener, TASKBAR_PINNING);
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java b/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java
index 1db3b5a..b877d7a 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java
@@ -16,8 +16,8 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.icons.ComponentWithLabelAndIcon;
-import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.BaseIconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 
 /**
@@ -26,8 +26,7 @@
  * (who's implementation is owned by the launcher). This object represents a widget type / class,
  * as opposed to a widget instance, and so should not be confused with {@link LauncherAppWidgetInfo}
  */
-public class LauncherAppWidgetProviderInfo extends AppWidgetProviderInfo
-        implements ComponentWithLabelAndIcon {
+public class LauncherAppWidgetProviderInfo extends AppWidgetProviderInfo implements CachedObject {
 
     public static final String CLS_CUSTOM_WIDGET_PREFIX = "#custom-widget-";
 
@@ -69,6 +68,8 @@
 
     protected boolean mIsMinSizeFulfilled;
 
+    private PackageManager mPM;
+
     public static LauncherAppWidgetProviderInfo fromProviderInfo(Context context,
             AppWidgetProviderInfo info) {
         final LauncherAppWidgetProviderInfo launcherInfo;
@@ -97,6 +98,7 @@
     }
 
     public void initSpans(Context context, InvariantDeviceProfile idp) {
+        mPM = context.getApplicationContext().getPackageManager();
         int minSpanX = 0;
         int minSpanY = 0;
         int maxSpanX = idp.numColumns;
@@ -104,7 +106,6 @@
         int spanX = 0;
         int spanY = 0;
 
-
         Point cellSize = new Point();
         for (DeviceProfile dp : idp.supportedProfiles) {
             dp.getCellSize(cellSize);
@@ -188,8 +189,9 @@
                 (widgetSize + widgetPadding + cellSpacing) / (cellSize + cellSpacing)));
     }
 
-    public String getLabel(PackageManager packageManager) {
-        return super.loadLabel(packageManager);
+    @Override
+    public CharSequence getLabel() {
+        return super.loadLabel(mPM);
     }
 
     public Point getMinSpans() {
@@ -225,7 +227,7 @@
     }
 
     @Override
-    public Drawable getFullResIcon(IconCache cache) {
+    public Drawable getFullResIcon(BaseIconCache cache) {
         return cache.getFullResIcon(getActivityInfo());
     }
 
diff --git a/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java b/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
index 5ad9222..82a6883 100644
--- a/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
+++ b/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
@@ -18,7 +18,6 @@
 
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.os.Parcel;
 import android.os.Parcelable;
 
@@ -64,7 +63,7 @@
     }
 
     @Override
-    public String getLabel(PackageManager packageManager) {
+    public CharSequence getLabel() {
         return Utilities.trim(label);
     }
 
diff --git a/src_no_quickstep/com/android/launcher3/util/StateManagerProtoLogProxy.java b/src_no_quickstep/com/android/launcher3/util/StateManagerProtoLogProxy.java
new file mode 100644
index 0000000..34e15f7
--- /dev/null
+++ b/src_no_quickstep/com/android/launcher3/util/StateManagerProtoLogProxy.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util;
+
+/**
+ * Proxy class used for StateManager ProtoLog support.
+ */
+public class StateManagerProtoLogProxy {
+
+    public static void logGoToState(Object fromState, Object toState, String trace) { }
+
+    public static void logCreateAtomicAnimation(Object fromState, Object toState, String trace) { }
+
+    public static void logOnStateTransitionStart(Object state) { }
+
+    public static void logOnStateTransitionEnd(Object state) { }
+
+    public static void logCancelAnimation(boolean animationOngoing, String trace) { }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/AppWidgetsRestoredReceiverTest.kt b/tests/multivalentTests/src/com/android/launcher3/AppWidgetsRestoredReceiverTest.kt
index 21abab4..0e06051 100644
--- a/tests/multivalentTests/src/com/android/launcher3/AppWidgetsRestoredReceiverTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/AppWidgetsRestoredReceiverTest.kt
@@ -29,7 +29,7 @@
 
     @Before
     fun setup() {
-        launcherPrefs = LauncherPrefs(DeviceHelpers.context)
+        launcherPrefs = LauncherPrefs.get(DeviceHelpers.context)
         receiverUnderTest = AppWidgetsRestoredReceiver()
     }
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
new file mode 100644
index 0000000..946bbc5
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3
+
+import android.content.Context
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+
+/** Emulates Launcher preferences for a test environment. */
+class FakeLauncherPrefs(private val context: Context) : LauncherPrefs() {
+    private val prefsMap = mutableMapOf<String, Any>()
+    private val listeners = mutableSetOf<LauncherPrefChangeListener>()
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <T> get(item: ContextualItem<T>): T {
+        return prefsMap.getOrDefault(item.sharedPrefKey, item.defaultValueFromContext(context)) as T
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <T> get(item: ConstantItem<T>): T {
+        return prefsMap.getOrDefault(item.sharedPrefKey, item.defaultValue) as T
+    }
+
+    override fun put(vararg itemsToValues: Pair<Item, Any>) = putSync(*itemsToValues)
+
+    override fun <T : Any> put(item: Item, value: T) = putSync(item to value)
+
+    override fun putSync(vararg itemsToValues: Pair<Item, Any>) {
+        itemsToValues
+            .map { (i, v) -> i.sharedPrefKey to v }
+            .forEach { (k, v) ->
+                prefsMap[k] = v
+                notifyChange(k)
+            }
+    }
+
+    override fun addListener(listener: LauncherPrefChangeListener, vararg items: Item) {
+        listeners.add(listener)
+    }
+
+    override fun removeListener(listener: LauncherPrefChangeListener, vararg items: Item) {
+        listeners.remove(listener)
+    }
+
+    override fun has(vararg items: Item) = items.all { it.sharedPrefKey in prefsMap }
+
+    override fun remove(vararg items: Item) = removeSync(*items)
+
+    override fun removeSync(vararg items: Item) {
+        items
+            .filter { it.sharedPrefKey in prefsMap }
+            .forEach {
+                prefsMap.remove(it.sharedPrefKey)
+                notifyChange(it.sharedPrefKey)
+            }
+    }
+
+    override fun close() = Unit
+
+    private fun notifyChange(key: String) {
+        // Mimics SharedPreferencesImpl#notifyListeners main thread dispatching.
+        MAIN_EXECUTOR.execute { listeners.forEach { it.onPrefChanged(key) } }
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt
new file mode 100644
index 0000000..2463c93
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt
@@ -0,0 +1,135 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3
+
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+private val TEST_CONSTANT_ITEM = LauncherPrefs.nonRestorableItem("TEST_BOOLEAN_ITEM", false)
+
+private val TEST_CONTEXTUAL_ITEM =
+    ContextualItem(
+        "TEST_CONTEXTUAL_ITEM",
+        true,
+        { false },
+        EncryptionType.ENCRYPTED,
+        Boolean::class.java,
+    )
+
+@RunWith(LauncherMultivalentJUnit::class)
+class FakeLauncherPrefsTest {
+    private val launcherPrefs = FakeLauncherPrefs(getApplicationContext())
+
+    @Test
+    fun testGet_constantItemNotInPrefs_returnsDefaultValue() {
+        assertThat(launcherPrefs.get(TEST_CONSTANT_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testGet_constantItemInPrefs_returnsStoredValue() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        assertThat(launcherPrefs.get(TEST_CONSTANT_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testGet_contextualItemNotInPrefs_returnsDefaultValue() {
+        assertThat(launcherPrefs.get(TEST_CONTEXTUAL_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testGet_contextualItemInPrefs_returnsStoredValue() {
+        launcherPrefs.put(TEST_CONTEXTUAL_ITEM, true)
+        assertThat(launcherPrefs.get(TEST_CONTEXTUAL_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testPut_multipleItems_storesAll() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM to true, TEST_CONTEXTUAL_ITEM to true)
+        assertThat(launcherPrefs.get(TEST_CONSTANT_ITEM)).isTrue()
+        assertThat(launcherPrefs.get(TEST_CONTEXTUAL_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testHas_itemNotInPrefs_returnsFalse() {
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testHas_itemInPrefs_returnsTrue() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testHas_twoItemsWithOneInPrefs_returnsFalse() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM, TEST_CONTEXTUAL_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testHas_twoItemsInPrefs_returnsTrue() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM to true, TEST_CONTEXTUAL_ITEM to true)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM, TEST_CONTEXTUAL_ITEM)).isTrue()
+    }
+
+    @Test
+    fun testRemove_itemInPrefs_removesItem() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        launcherPrefs.remove(TEST_CONSTANT_ITEM)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testRemove_itemsInPrefs_removesItems() {
+        launcherPrefs.put(TEST_CONSTANT_ITEM to true, TEST_CONTEXTUAL_ITEM to true)
+        launcherPrefs.remove(TEST_CONSTANT_ITEM, TEST_CONTEXTUAL_ITEM)
+        assertThat(launcherPrefs.has(TEST_CONSTANT_ITEM, TEST_CONTEXTUAL_ITEM)).isFalse()
+    }
+
+    @Test
+    fun testAddListener_changeItemInPrefs_callsListener() {
+        var changedKey: String? = null
+        launcherPrefs.addListener({ changedKey = it }, TEST_CONSTANT_ITEM)
+        getInstrumentation().runOnMainSync { launcherPrefs.put(TEST_CONSTANT_ITEM, true) }
+        assertThat(changedKey).isEqualTo(TEST_CONSTANT_ITEM.sharedPrefKey)
+    }
+
+    @Test
+    fun testAddListener_removeItemFromPrefs_callsListener() {
+        var changedKey: String? = null
+        launcherPrefs.put(TEST_CONSTANT_ITEM, true)
+        launcherPrefs.addListener({ changedKey = it }, TEST_CONSTANT_ITEM)
+
+        getInstrumentation().runOnMainSync { launcherPrefs.remove(TEST_CONSTANT_ITEM) }
+        assertThat(changedKey).isEqualTo(TEST_CONSTANT_ITEM.sharedPrefKey)
+    }
+
+    @Test
+    fun testRemoveListener_changeItemInPrefs_doesNotCallListener() {
+        var changedKey: String? = null
+        val listener = LauncherPrefChangeListener { changedKey = it }
+        launcherPrefs.addListener(listener, TEST_CONSTANT_ITEM)
+
+        launcherPrefs.removeListener(listener)
+        getInstrumentation().runOnMainSync { launcherPrefs.put(TEST_CONSTANT_ITEM, true) }
+        assertThat(changedKey).isNull()
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/LauncherPrefsTest.kt b/tests/multivalentTests/src/com/android/launcher3/LauncherPrefsTest.kt
index b813095..4aeef2e 100644
--- a/tests/multivalentTests/src/com/android/launcher3/LauncherPrefsTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/LauncherPrefsTest.kt
@@ -17,7 +17,6 @@
 
 import android.content.Context
 import android.content.SharedPreferences
-import android.content.SharedPreferences.OnSharedPreferenceChangeListener
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import androidx.test.platform.app.InstrumentationRegistry
@@ -63,7 +62,7 @@
     @Test
     fun addListener_listeningForStringItemUpdates_isCorrectlyNotifiedOfUpdates() {
         val latch = CountDownLatch(1)
-        val listener = OnSharedPreferenceChangeListener { _, _ -> latch.countDown() }
+        val listener = LauncherPrefChangeListener { latch.countDown() }
 
         with(launcherPrefs) {
             putSync(TEST_STRING_ITEM.to(TEST_STRING_ITEM.defaultValue))
@@ -78,7 +77,7 @@
     @Test
     fun removeListener_previouslyListeningForStringItemUpdates_isNoLongerNotifiedOfUpdates() {
         val latch = CountDownLatch(1)
-        val listener = OnSharedPreferenceChangeListener { _, _ -> latch.countDown() }
+        val listener = LauncherPrefChangeListener { latch.countDown() }
 
         with(launcherPrefs) {
             addListener(listener, TEST_STRING_ITEM)
@@ -94,14 +93,14 @@
     @Test
     fun addListenerAndRemoveListener_forMultipleItems_bothWorkProperly() {
         var latch = CountDownLatch(3)
-        val listener = OnSharedPreferenceChangeListener { _, _ -> latch.countDown() }
+        val listener = LauncherPrefChangeListener { latch.countDown() }
 
         with(launcherPrefs) {
             addListener(listener, TEST_INT_ITEM, TEST_STRING_ITEM, TEST_BOOLEAN_ITEM)
             putSync(
                 TEST_INT_ITEM.to(TEST_INT_ITEM.defaultValue + 123),
                 TEST_STRING_ITEM.to(TEST_STRING_ITEM.defaultValue + "abc"),
-                TEST_BOOLEAN_ITEM.to(!TEST_BOOLEAN_ITEM.defaultValue)
+                TEST_BOOLEAN_ITEM.to(!TEST_BOOLEAN_ITEM.defaultValue),
             )
             assertThat(latch.await(WAIT_TIME_IN_SECONDS, TimeUnit.SECONDS)).isTrue()
 
@@ -110,7 +109,7 @@
             putSync(
                 TEST_INT_ITEM.to(TEST_INT_ITEM.defaultValue),
                 TEST_STRING_ITEM.to(TEST_STRING_ITEM.defaultValue),
-                TEST_BOOLEAN_ITEM.to(TEST_BOOLEAN_ITEM.defaultValue)
+                TEST_BOOLEAN_ITEM.to(TEST_BOOLEAN_ITEM.defaultValue),
             )
             remove(TEST_INT_ITEM, TEST_STRING_ITEM, TEST_BOOLEAN_ITEM)
 
@@ -150,7 +149,7 @@
             putSync(
                 TEST_STRING_ITEM.to(TEST_STRING_ITEM.defaultValue),
                 TEST_INT_ITEM.to(TEST_INT_ITEM.defaultValue),
-                TEST_BOOLEAN_ITEM.to(TEST_BOOLEAN_ITEM.defaultValue)
+                TEST_BOOLEAN_ITEM.to(TEST_BOOLEAN_ITEM.defaultValue),
             )
             assertThat(has(TEST_BOOLEAN_ITEM, TEST_INT_ITEM, TEST_STRING_ITEM)).isTrue()
             remove(TEST_STRING_ITEM, TEST_INT_ITEM, TEST_BOOLEAN_ITEM)
@@ -191,7 +190,7 @@
             LauncherPrefs.backedUpItem(
                 TEST_PREF_KEY,
                 TEST_DEFAULT_VALUE,
-                EncryptionType.DEVICE_PROTECTED
+                EncryptionType.DEVICE_PROTECTED,
             )
 
         val bootAwarePrefs: SharedPreferences =
@@ -212,7 +211,7 @@
             LauncherPrefs.backedUpItem(
                 TEST_PREF_KEY,
                 TEST_DEFAULT_VALUE,
-                EncryptionType.DEVICE_PROTECTED
+                EncryptionType.DEVICE_PROTECTED,
             )
 
         val bootAwarePrefs: SharedPreferences =
diff --git a/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt b/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt
index 8e54c94..ed9a080 100644
--- a/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt
@@ -18,10 +18,8 @@
 
 import android.content.ComponentName
 import android.content.pm.ApplicationInfo
-import android.content.pm.PackageManager
 import android.database.MatrixCursor
 import android.os.Process.myUserHandle
-import androidx.test.core.app.ApplicationProvider.getApplicationContext
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.icons.cache.BaseIconCache
@@ -49,7 +47,7 @@
     @Mock private lateinit var baseIconCache: BaseIconCache
 
     private var cursor: MatrixCursor? = null
-    private var cachingLogic = CachedObjectCachingLogic<BaseIconCache>(getApplicationContext())
+    private var cachingLogic = CachedObjectCachingLogic
 
     @Before
     fun setup() {
@@ -137,14 +135,13 @@
     }
 }
 
-class TestCachedObject(val cn: ComponentName, val freshnessId: String) :
-    CachedObject<BaseIconCache> {
+class TestCachedObject(val cn: ComponentName, val freshnessId: String) : CachedObject {
 
     override fun getComponent() = cn
 
     override fun getUser() = myUserHandle()
 
-    override fun getLabel(pm: PackageManager?): CharSequence? = null
+    override fun getLabel(): CharSequence? = null
 
     override fun getApplicationInfo(): ApplicationInfo? = null
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/AsyncBindingTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/AsyncBindingTest.kt
index c91577f..ba59253 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/AsyncBindingTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/AsyncBindingTest.kt
@@ -136,7 +136,7 @@
     @Test
     fun test_bind_sync_partially_inflates_on_background() {
         modelHelper.loadModelSync()
-        assertTrue(modelHelper.model.isModelLoaded)
+        assertTrue(modelHelper.model.isModelLoaded())
         callbacks.inflater = itemInflater
 
         val firstPageBindIds = IntSet()
diff --git a/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index b3675a6..d0c168a 100644
--- a/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -101,7 +101,7 @@
         mMockController = Mockito.mock(ModelDbController.class);
         mMockDb = mock(SQLiteDatabase.class);
         mMockCursor = mock(Cursor.class);
-        mPrefs = new LauncherPrefs(mContext);
+        mPrefs = LauncherPrefs.get(mContext);
         mMockRestoreEventLogger = mock(LauncherRestoreEventLogger.class);
     }
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt
index 0a3035a..af2c378 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt
@@ -17,7 +17,6 @@
 package com.android.launcher3.widget.custom
 
 import android.content.ComponentName
-import android.content.pm.PackageManager
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
@@ -25,7 +24,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mockito.mock
 
 @MediumTest
 @RunWith(AndroidJUnit4::class)
@@ -47,7 +45,7 @@
     @Test
     fun get_label() {
         underTest.label = "  TEST_LABEL"
-        assertEquals(LABEL_NAME, underTest.getLabel(mock(PackageManager::class.java)))
+        assertEquals(LABEL_NAME, underTest.getLabel())
     }
 
     companion object {
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/model/WidgetsListBaseEntriesBuilderTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/model/WidgetsListBaseEntriesBuilderTest.kt
index 5df7caa..063ab32 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/model/WidgetsListBaseEntriesBuilderTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/model/WidgetsListBaseEntriesBuilderTest.kt
@@ -26,8 +26,8 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherAppState
-import com.android.launcher3.icons.ComponentWithLabel
 import com.android.launcher3.icons.IconCache
+import com.android.launcher3.icons.cache.CachedObject
 import com.android.launcher3.model.WidgetItem
 import com.android.launcher3.model.data.PackageItemInfo
 import com.android.launcher3.util.ActivityContextWrapper
@@ -66,11 +66,11 @@
         testInvariantProfile = LauncherAppState.getIDP(context)
 
         doAnswer { invocation: InvocationOnMock ->
-                val componentWithLabel = invocation.getArgument<Any>(0) as ComponentWithLabel
+                val componentWithLabel = invocation.getArgument<Any>(0) as CachedObject
                 componentWithLabel.getComponent().shortClassName
             }
             .`when`(iconCache)
-            .getTitleNoCache(any<ComponentWithLabel>())
+            .getTitleNoCache(any<CachedObject>())
         underTest = WidgetsListBaseEntriesBuilder(context)
 
         allWidgets =
@@ -79,14 +79,14 @@
                 packageItemInfoWithTitle(APP_1_PACKAGE_NAME, APP_1_PACKAGE_TITLE) to
                     listOf(
                         createWidgetItem(APP_1_PACKAGE_NAME, APP_1_PROVIDER_1_CLASS_NAME),
-                        createWidgetItem(APP_1_PACKAGE_NAME, APP_1_PROVIDER_2_CLASS_NAME)
+                        createWidgetItem(APP_1_PACKAGE_NAME, APP_1_PROVIDER_2_CLASS_NAME),
                     ),
                 // app 2
                 packageItemInfoWithTitle(APP_2_PACKAGE_NAME, APP_2_PACKAGE_TITLE) to
                     listOf(createWidgetItem(APP_2_PACKAGE_NAME, APP_2_PROVIDER_1_CLASS_NAME)),
                 // app 3
                 packageItemInfoWithTitle(APP_3_PACKAGE_NAME, APP_3_PACKAGE_TITLE) to
-                    listOf(createWidgetItem(APP_3_PACKAGE_NAME, APP_3_PROVIDER_1_CLASS_NAME))
+                    listOf(createWidgetItem(APP_3_PACKAGE_NAME, APP_3_PROVIDER_1_CLASS_NAME)),
             )
     }
 
@@ -96,7 +96,7 @@
             listOf(
                 APP_1_EXPECTED_SECTION_NAME to 2,
                 APP_2_EXPECTED_SECTION_NAME to 1,
-                APP_3_EXPECTED_SECTION_NAME to 1
+                APP_3_EXPECTED_SECTION_NAME to 1,
             )
 
         val entries = underTest.build(allWidgets)
@@ -122,7 +122,7 @@
         val expectedWidgetsCountBySection =
             listOf(
                 APP_1_EXPECTED_SECTION_NAME to 1, // one widget filtered out
-                APP_3_EXPECTED_SECTION_NAME to 1
+                APP_3_EXPECTED_SECTION_NAME to 1,
             )
 
         val entries =
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
index d4e061a..c9b6d4f 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -42,8 +42,8 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
@@ -87,7 +87,7 @@
         mTestProfile.numColumns = 5;
 
         doAnswer(invocation -> {
-            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            CachedObject componentWithLabel = invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
         mViewHolderBinder = new WidgetsListHeaderViewHolderBinder(
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
index e1cc010..0d9464a 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
@@ -45,8 +45,8 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
@@ -92,7 +92,7 @@
         mTestProfile.numColumns = 5;
 
         doAnswer(invocation -> {
-            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            CachedObject componentWithLabel = invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt
index 1822639..1da74cb 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetPickerDataProviderTest.kt
@@ -27,8 +27,8 @@
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.LauncherSettings
-import com.android.launcher3.icons.ComponentWithLabel
 import com.android.launcher3.icons.IconCache
+import com.android.launcher3.icons.cache.CachedObject
 import com.android.launcher3.model.WidgetItem
 import com.android.launcher3.model.data.PackageItemInfo
 import com.android.launcher3.util.ActivityContextWrapper
@@ -81,11 +81,11 @@
         testInvariantProfile = LauncherAppState.getIDP(context)
 
         doAnswer { invocation: InvocationOnMock ->
-                val componentWithLabel = invocation.getArgument<Any>(0) as ComponentWithLabel
+                val componentWithLabel = invocation.getArgument<Any>(0) as CachedObject
                 componentWithLabel.getComponent().shortClassName
             }
             .`when`(iconCache)
-            .getTitleNoCache(any<ComponentWithLabel>())
+            .getTitleNoCache(any<CachedObject>())
 
         appWidgetItem = createWidgetItem()
     }
@@ -113,8 +113,8 @@
             listOf(
                 PendingAddWidgetInfo(
                     appWidgetItem.widgetInfo,
-                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
-                ),
+                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION,
+                )
             )
         underTest.setWidgetRecommendations(recommendations)
 
@@ -133,8 +133,8 @@
             listOf(
                 PendingAddWidgetInfo(
                     appWidgetItem.widgetInfo,
-                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
-                ),
+                    LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION,
+                )
             )
         underTest.setWidgetRecommendations(recommendations)
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
index 7552619..6088c8e 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
@@ -33,8 +33,8 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
@@ -81,7 +81,7 @@
         mTestProfile.numColumns = 5;
 
         doAnswer(invocation -> {
-            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            CachedObject componentWithLabel = invocation.getArgument(0);
             return mWidgetsToLabels.get(componentWithLabel.getComponent());
         }).when(mIconCache).getTitleNoCache(any());
     }
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt
index e59e211..deec67a 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/data/WidgetPickerDataTest.kt
@@ -27,8 +27,8 @@
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION
-import com.android.launcher3.icons.ComponentWithLabel
 import com.android.launcher3.icons.IconCache
+import com.android.launcher3.icons.cache.CachedObject
 import com.android.launcher3.model.WidgetItem
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.model.data.PackageItemInfo
@@ -86,11 +86,11 @@
         testInvariantProfile = LauncherAppState.getIDP(context)
 
         doAnswer { invocation: InvocationOnMock ->
-                val componentWithLabel = invocation.getArgument<Any>(0) as ComponentWithLabel
+                val componentWithLabel = invocation.getArgument<Any>(0) as CachedObject
                 componentWithLabel.getComponent().shortClassName
             }
             .`when`(iconCache)
-            .getTitleNoCache(any<ComponentWithLabel>())
+            .getTitleNoCache(any<CachedObject>())
 
         app1PackageItemInfo = packageItemInfoWithTitle(APP_1_PACKAGE_NAME, APP_1_PACKAGE_TITLE)
         app2PackageItemInfo = packageItemInfoWithTitle(APP_2_PACKAGE_NAME, APP_2_PACKAGE_TITLE)
@@ -123,7 +123,7 @@
         val widgetPickerData =
             WidgetPickerData(
                 allWidgets = appTwoWidgetsListBaseEntries(),
-                defaultWidgets = appTwoWidgetsListBaseEntries()
+                defaultWidgets = appTwoWidgetsListBaseEntries(),
             )
 
         val newWidgetData =
@@ -143,19 +143,19 @@
                         addAll(appOneWidgetsListBaseEntries())
                         addAll(appTwoWidgetsListBaseEntries())
                     },
-                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() }
+                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() },
             )
         val recommendations: List<ItemInfo> =
             listOf(
                 PendingAddWidgetInfo(
                     app1WidgetItem1.widgetInfo,
                     CONTAINER_WIDGETS_PREDICTION,
-                    CATEGORY_1
+                    CATEGORY_1,
                 ),
                 PendingAddWidgetInfo(
                     app2WidgetItem1.widgetInfo,
                     CONTAINER_WIDGETS_PREDICTION,
-                    CATEGORY_2
+                    CATEGORY_2,
                 ),
             )
 
@@ -175,7 +175,7 @@
                         addAll(appOneWidgetsListBaseEntries())
                         addAll(appTwoWidgetsListBaseEntries())
                     },
-                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() }
+                defaultWidgets = buildList { appTwoWidgetsListBaseEntries() },
             )
         val recommendations: List<ItemInfo> =
             listOf(
@@ -201,7 +201,7 @@
                         addAll(appTwoWidgetsListBaseEntries())
                     },
                 defaultWidgets = buildList { appTwoWidgetsListBaseEntries() },
-                recommendations = mapOf(CATEGORY_1 to listOf(app1WidgetItem1))
+                recommendations = mapOf(CATEGORY_1 to listOf(app1WidgetItem1)),
             )
 
         val updatedData = widgetPickerData.withRecommendedWidgets(listOf())
@@ -242,7 +242,7 @@
                         addAll(appOneWidgetsListBaseEntries())
                         addAll(appTwoWidgetsListBaseEntries())
                     },
-                defaultWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) }
+                defaultWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) },
             )
         val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
 
@@ -263,7 +263,7 @@
                         addAll(appTwoWidgetsListBaseEntries())
                     },
                 defaultWidgets =
-                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) }
+                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) },
             )
         val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
 
@@ -271,7 +271,7 @@
             findContentEntryForPackageUser(
                 widgetPickerData = widgetPickerData,
                 packageUserKey = app1PackageUserKey,
-                fromDefaultWidgets = true
+                fromDefaultWidgets = true,
             )
 
         assertThat(contentEntry).isNotNull()
@@ -302,7 +302,7 @@
                         addAll(appTwoWidgetsListBaseEntries())
                     },
                 defaultWidgets =
-                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) }
+                    buildList { addAll(appOneWidgetsListBaseEntries(includeWidgetTwo = false)) },
             )
 
         val widgets = findAllWidgetsForPackageUser(widgetPickerData, app1PackageUserKey)
@@ -314,9 +314,7 @@
     @Test
     fun findAllWidgetsForPackageUser_noMatch_returnsEmptyList() {
         val widgetPickerData =
-            WidgetPickerData(
-                allWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) },
-            )
+            WidgetPickerData(allWidgets = buildList { addAll(appTwoWidgetsListBaseEntries()) })
         val app1PackageUserKey = PackageUserKey.fromPackageItemInfo(app1PackageItemInfo)
 
         val widgets = findAllWidgetsForPackageUser(widgetPickerData, app1PackageUserKey)
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
index 24d66a3..59f352b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
@@ -41,8 +41,8 @@
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.icons.BitmapInfo;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.search.SearchCallback;
@@ -87,7 +87,7 @@
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         doAnswer(invocation -> {
-            ComponentWithLabel componentWithLabel = (ComponentWithLabel) invocation.getArgument(0);
+            CachedObject componentWithLabel = invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
         mTestProfile = new InvariantDeviceProfile();
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
index 7adb2b1..2f5fcfe 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
@@ -28,7 +28,6 @@
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.graphics.Point;
 import android.graphics.drawable.Drawable;
 import android.os.UserHandle;
@@ -39,8 +38,9 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
+import com.android.launcher3.icons.cache.BaseIconCache;
+import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
@@ -99,7 +99,7 @@
         initTestWidgets();
         initTestShortcuts();
 
-        doAnswer(invocation -> ((ComponentWithLabel) invocation.getArgument(0))
+        doAnswer(invocation -> ((CachedObject) invocation.getArgument(0))
                 .getComponent().getPackageName())
                 .when(mIconCache).getTitleNoCache(any());
     }
@@ -280,16 +280,15 @@
     }
 
     private void initTestShortcuts() {
-        PackageManager packageManager = mContext.getPackageManager();
         mShortcut1 = new WidgetItem(new TestShortcutConfigActivityInfo(
                 ComponentName.createRelative(TEST_PACKAGE, ".shortcut1"), UserHandle.CURRENT),
-                mIconCache, packageManager);
+                mIconCache);
         mShortcut2 = new WidgetItem(new TestShortcutConfigActivityInfo(
                 ComponentName.createRelative(TEST_PACKAGE, ".shortcut2"), UserHandle.CURRENT),
-                mIconCache, packageManager);
+                mIconCache);
         mShortcut3 = new WidgetItem(new TestShortcutConfigActivityInfo(
                 ComponentName.createRelative(TEST_PACKAGE, ".shortcut3"), UserHandle.CURRENT),
-                mIconCache, packageManager);
+                mIconCache);
 
     }
 
@@ -300,12 +299,12 @@
         }
 
         @Override
-        public Drawable getFullResIcon(IconCache cache) {
+        public Drawable getFullResIcon(BaseIconCache cache) {
             return null;
         }
 
         @Override
-        public CharSequence getLabel(PackageManager pm) {
+        public CharSequence getLabel() {
             return null;
         }
     }
diff --git a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
index 0dd13a9..b17cd4d 100644
--- a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
@@ -33,17 +33,15 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
-import org.mockito.ArgumentMatchers.any
-import org.mockito.ArgumentMatchers.anyInt
-import org.mockito.ArgumentMatchers.anyList
-import org.mockito.ArgumentMatchers.anyMap
 import org.mockito.Mock
 import org.mockito.Mockito
+import org.mockito.Mockito.mock
 import org.mockito.Mockito.times
 import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations
 import org.mockito.MockitoSession
 import org.mockito.Spy
+import org.mockito.kotlin.any
 import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.spy
@@ -67,7 +65,7 @@
             installedHotseatItems = mutableSetOf("installedHotseatItem"),
             installedWorkspaceItems = mutableSetOf("installedWorkspaceItem"),
             firstScreenInstalledWidgets = mutableSetOf("installedFirstScreenWidget"),
-            secondaryScreenInstalledWidgets = mutableSetOf("installedSecondaryScreenWidget")
+            secondaryScreenInstalledWidgets = mutableSetOf("installedSecondaryScreenWidget"),
         )
     private lateinit var mockitoSession: MockitoSession
 
@@ -75,7 +73,7 @@
     @Mock private lateinit var bgAllAppsList: AllAppsList
     @Mock private lateinit var modelDelegate: ModelDelegate
     @Mock private lateinit var launcherBinder: BaseLauncherBinder
-    @Mock private lateinit var launcherModel: LauncherModel
+    private lateinit var launcherModel: LauncherModel
     @Mock private lateinit var transaction: LoaderTransaction
     @Mock private lateinit var iconCache: IconCache
     @Mock private lateinit var idleLock: LooperIdleLock
@@ -89,6 +87,7 @@
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
+        launcherModel = mock(LauncherModel::class.java)
         mockitoSession =
             ExtendedMockito.mockitoSession()
                 .strictness(Strictness.LENIENT)
@@ -105,15 +104,16 @@
 
         doReturn(TestViewHelpers.findWidgetProvider(false))
             .`when`(context.spyService(AppWidgetManager::class.java))
-            .getAppWidgetInfo(anyInt())
+            .getAppWidgetInfo(any())
         `when`(app.context).thenReturn(context)
         `when`(app.model).thenReturn(launcherModel)
-        `when`(launcherModel.beginLoader(any(LoaderTask::class.java))).thenReturn(transaction)
+
+        `when`(launcherModel.beginLoader(any())).thenReturn(transaction)
         `when`(app.iconCache).thenReturn(iconCache)
         `when`(launcherModel.modelDbController)
             .thenReturn(FactitiousDbController(context, INSERTION_STATEMENT_FILE))
         `when`(app.invariantDeviceProfile).thenReturn(idp)
-        `when`(launcherBinder.newIdleLock(any(LoaderTask::class.java))).thenReturn(idleLock)
+        `when`(launcherBinder.newIdleLock(any())).thenReturn(idleLock)
         `when`(idleLock.awaitLocked(1000)).thenReturn(false)
         `when`(iconCache.updateHandler).thenReturn(iconCacheUpdateHandler)
         context.putObject(UserCache.INSTANCE, userCache)
@@ -149,12 +149,12 @@
 
         verify(launcherBinder).bindWorkspace(true, false)
         verify(modelDelegate).workspaceLoadComplete()
-        verify(modelDelegate).loadAndBindAllAppsItems(any(), any(), any())
+        verify(modelDelegate).loadAndBindAllAppsItems(any(), anyOrNull(), any())
         verify(launcherBinder).bindAllApps()
         verify(iconCacheUpdateHandler, times(4)).updateIcons(any(), any<CachingLogic<Any>>(), any())
         verify(launcherBinder).bindDeepShortcuts()
         verify(launcherBinder).bindWidgets()
-        verify(modelDelegate).loadAndBindOtherItems(any())
+        verify(modelDelegate).loadAndBindOtherItems(anyOrNull())
         verify(iconCacheUpdateHandler).finish()
         verify(modelDelegate).modelLoadComplete()
         verify(transaction).commit()
@@ -209,10 +209,10 @@
         `when`(app.context).thenReturn(spyContext)
         whenever(
                 FirstScreenBroadcastHelper.createModelsForFirstScreenBroadcast(
-                    anyOrNull(),
-                    anyList(),
-                    anyMap(),
-                    anyList()
+                    any(),
+                    any(),
+                    any(),
+                    any(),
                 )
             )
             .thenReturn(listOf(expectedBroadcastModel))
@@ -220,7 +220,7 @@
         whenever(
                 FirstScreenBroadcastHelper.sendBroadcastsForModels(
                     spyContext,
-                    listOf(expectedBroadcastModel)
+                    listOf(expectedBroadcastModel),
                 )
             )
             .thenCallRealMethod()
@@ -239,34 +239,34 @@
         assertEquals(expectedBroadcastModel.installerPackage, actualBroadcastIntent.`package`)
         assertEquals(
             ArrayList(expectedBroadcastModel.installedWorkspaceItems),
-            actualBroadcastIntent.getStringArrayListExtra("workspaceInstalledItems")
+            actualBroadcastIntent.getStringArrayListExtra("workspaceInstalledItems"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.installedHotseatItems),
-            actualBroadcastIntent.getStringArrayListExtra("hotseatInstalledItems")
+            actualBroadcastIntent.getStringArrayListExtra("hotseatInstalledItems"),
         )
         assertEquals(
             ArrayList(
                 expectedBroadcastModel.firstScreenInstalledWidgets +
                     expectedBroadcastModel.secondaryScreenInstalledWidgets
             ),
-            actualBroadcastIntent.getStringArrayListExtra("widgetInstalledItems")
+            actualBroadcastIntent.getStringArrayListExtra("widgetInstalledItems"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.pendingCollectionItems),
-            actualBroadcastIntent.getStringArrayListExtra("folderItem")
+            actualBroadcastIntent.getStringArrayListExtra("folderItem"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.pendingWorkspaceItems),
-            actualBroadcastIntent.getStringArrayListExtra("workspaceItem")
+            actualBroadcastIntent.getStringArrayListExtra("workspaceItem"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.pendingHotseatItems),
-            actualBroadcastIntent.getStringArrayListExtra("hotseatItem")
+            actualBroadcastIntent.getStringArrayListExtra("hotseatItem"),
         )
         assertEquals(
             ArrayList(expectedBroadcastModel.pendingWidgetItems),
-            actualBroadcastIntent.getStringArrayListExtra("widgetItem")
+            actualBroadcastIntent.getStringArrayListExtra("widgetItem"),
         )
     }
 
@@ -277,10 +277,10 @@
         `when`(app.context).thenReturn(spyContext)
         whenever(
                 FirstScreenBroadcastHelper.createModelsForFirstScreenBroadcast(
-                    anyOrNull(),
-                    anyList(),
-                    anyMap(),
-                    anyList()
+                    any(),
+                    any(),
+                    any(),
+                    any(),
                 )
             )
             .thenReturn(listOf(expectedBroadcastModel))
@@ -288,7 +288,7 @@
         whenever(
                 FirstScreenBroadcastHelper.sendBroadcastsForModels(
                     spyContext,
-                    listOf(expectedBroadcastModel)
+                    listOf(expectedBroadcastModel),
                 )
             )
             .thenCallRealMethod()
@@ -300,7 +300,7 @@
             .runSyncOnBackgroundThread()
 
         // Then
-        verify(spyContext, times(0)).sendBroadcast(any(Intent::class.java))
+        verify(spyContext, times(0)).sendBroadcast(any())
     }
 
     @Test
@@ -310,10 +310,10 @@
         `when`(app.context).thenReturn(spyContext)
         whenever(
                 FirstScreenBroadcastHelper.createModelsForFirstScreenBroadcast(
-                    anyOrNull(),
-                    anyList(),
-                    anyMap(),
-                    anyList()
+                    any(),
+                    any(),
+                    any(),
+                    any(),
                 )
             )
             .thenReturn(listOf(expectedBroadcastModel))
@@ -321,7 +321,7 @@
         whenever(
                 FirstScreenBroadcastHelper.sendBroadcastsForModels(
                     spyContext,
-                    listOf(expectedBroadcastModel)
+                    listOf(expectedBroadcastModel),
                 )
             )
             .thenCallRealMethod()
@@ -334,7 +334,7 @@
             .runSyncOnBackgroundThread()
 
         // Then
-        verify(spyContext, times(0)).sendBroadcast(any(Intent::class.java))
+        verify(spyContext, times(0)).sendBroadcast(any())
     }
 }
 
diff --git a/tests/src/com/android/launcher3/pm/InstallSessionHelperTest.kt b/tests/src/com/android/launcher3/pm/InstallSessionHelperTest.kt
index a991981..ca2ef42 100644
--- a/tests/src/com/android/launcher3/pm/InstallSessionHelperTest.kt
+++ b/tests/src/com/android/launcher3/pm/InstallSessionHelperTest.kt
@@ -28,7 +28,6 @@
 import androidx.test.filters.SmallTest
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.LauncherPrefs.Companion.PROMISE_ICON_IDS
-import com.android.launcher3.util.DaggerSingletonTracker
 import com.android.launcher3.util.Executors.MODEL_EXECUTOR
 import com.android.launcher3.util.IntArray
 import com.android.launcher3.util.LauncherModelHelper
@@ -54,7 +53,6 @@
     private val expectedAppPackage = "expectedAppPackage"
     private val expectedInstallerPackage = "expectedInstallerPackage"
     private val mockPackageInstaller: PackageInstaller = mock()
-    private val mTracker: DaggerSingletonTracker = mock()
 
     private lateinit var installSessionHelper: InstallSessionHelper
     private lateinit var launcherApps: LauncherApps
@@ -64,7 +62,7 @@
         whenever(packageManager.packageInstaller).thenReturn(mockPackageInstaller)
         whenever(sandboxContext.packageName).thenReturn(expectedInstallerPackage)
         launcherApps = sandboxContext.spyService(LauncherApps::class.java)
-        installSessionHelper = InstallSessionHelper(sandboxContext, mTracker)
+        installSessionHelper = InstallSessionHelper(sandboxContext)
     }
 
     @Test
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
index e6e02b4..961e7fc 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
@@ -93,7 +93,7 @@
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
         mLauncher.getWorkspace()
                 .openAllWidgets()
-                .getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .getWidget(mWidgetInfo.getLabel())
                 .dragToWorkspace(true, false);
         // Widget id for which the config activity was opened
         mWidgetId = monitor.getWidgetId();
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
index 9c916fa..9a2147a 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
@@ -61,7 +61,7 @@
         WidgetResizeFrame resizeFrame = mLauncher
                 .getWorkspace()
                 .openAllWidgets()
-                .getWidget(widgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .getWidget(widgetInfo.getLabel())
                 .dragWidgetToWorkspace();
 
         assertNotNull("Widget resize frame not shown after widget add", resizeFrame);
@@ -111,7 +111,7 @@
         WidgetResizeFrame resizeFrame = mLauncher
                 .getWorkspace()
                 .openAllWidgets()
-                .getWidget(widgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .getWidget(widgetInfo.getLabel())
                 .dragWidgetToWorkspace();
 
         assertNotNull("Widget resize frame not shown after widget add", resizeFrame);
diff --git a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
index a148744..d653317 100644
--- a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
@@ -16,8 +16,6 @@
 package com.android.launcher3.ui.workspace;
 
 import static com.android.launcher3.util.TestConstants.AppNames.TEST_APP_NAME;
-import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
-import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -40,8 +38,6 @@
 import com.android.launcher3.tapl.HomeAppIconMenuItem;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.Executors;
-import com.android.launcher3.util.rule.ScreenRecordRule;
-import com.android.launcher3.util.rule.TestStabilityRule;
 
 import org.junit.Test;
 
@@ -115,8 +111,6 @@
     }
 
     @Test
-    @TestStabilityRule.Stability(flavors = LOCAL | PLATFORM_POSTSUBMIT) // b/350557998
-    @ScreenRecordRule.ScreenRecord // b/350557998
     public void testShortcutIconWithTheme() throws Exception {
         setThemeEnabled(true);
         initialize(this);
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 6387b05..3097d9c 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -116,8 +116,8 @@
     }
 
     /** Get widget with supplied text. */
-    public Widget getWidget(String labelText) {
-        return getWidget(labelText, null);
+    public Widget getWidget(CharSequence labelText) {
+        return getWidget(labelText.toString(), null);
     }
 
     /** Get widget with supplied text and app package */