diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 3654e5f..abb763a 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -376,6 +376,7 @@
                             (flags & FLAG_ONLY_BACK_FOR_BOUNCER_VISIBLE) != 0 ||
                             (flags & FLAG_KEYGUARD_OCCLUDED) != 0;
                     return (flags & FLAG_DISABLE_BACK) == 0
+                            && (!mContext.isGestureNav() || !mContext.isUserSetupComplete())
                             && ((flags & FLAG_KEYGUARD_VISIBLE) == 0 || showingOnKeyguard);
                 }));
         // Hide back button in SUW if keyboard is showing (IME draws its own back).
@@ -406,8 +407,8 @@
         mPropertyHolders.add(
                 new StatePropertyHolder(mHomeButtonAlpha.get(
                         ALPHA_INDEX_KEYGUARD_OR_DISABLE),
-                flags -> (flags & FLAG_KEYGUARD_VISIBLE) == 0 &&
-                        (flags & FLAG_DISABLE_HOME) == 0));
+                        flags -> (flags & FLAG_KEYGUARD_VISIBLE) == 0
+                                && (flags & FLAG_DISABLE_HOME) == 0 && !mContext.isGestureNav()));
 
         // Recents button
         mRecentsButton = addButton(R.drawable.ic_sysbar_recent, BUTTON_RECENTS,
@@ -425,7 +426,7 @@
         });
         mPropertyHolders.add(new StatePropertyHolder(mRecentsButton,
                 flags -> (flags & FLAG_KEYGUARD_VISIBLE) == 0 && (flags & FLAG_DISABLE_RECENTS) == 0
-                        && !mContext.isNavBarKidsModeActive()));
+                        && !mContext.isNavBarKidsModeActive() && !mContext.isGestureNav()));
 
         // A11y button
         mA11yButton = addButton(R.drawable.ic_sysbar_accessibility_button, BUTTON_A11Y,
diff --git a/quickstep/src/com/android/launcher3/taskbar/StashedHandleView.java b/quickstep/src/com/android/launcher3/taskbar/StashedHandleView.java
index 83e4571..94e2244 100644
--- a/quickstep/src/com/android/launcher3/taskbar/StashedHandleView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/StashedHandleView.java
@@ -15,9 +15,12 @@
  */
 package com.android.launcher3.taskbar;
 
+import static com.android.launcher3.LauncherAnimUtils.SCALE_PROPERTY;
+
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ObjectAnimator;
+import android.animation.PropertyValuesHolder;
 import android.annotation.Nullable;
 import android.content.Context;
 import android.graphics.Rect;
@@ -111,4 +114,17 @@
             setBackgroundColor(newColor);
         }
     }
+
+    /**
+     * Updates the handle scale.
+     *
+     * @param scale target scale to animate towards (starting from current scale)
+     * @param durationMs milliseconds for the animation to take
+     */
+    public void animateScale(float scale, long durationMs) {
+        ObjectAnimator scaleAnim = ObjectAnimator.ofPropertyValuesHolder(this,
+                PropertyValuesHolder.ofFloat(SCALE_PROPERTY, scale));
+        scaleAnim.setDuration(durationMs).setAutoCancel(true);
+        scaleAnim.start();
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/StashedHandleViewController.java b/quickstep/src/com/android/launcher3/taskbar/StashedHandleViewController.java
index f258b47..36e054a 100644
--- a/quickstep/src/com/android/launcher3/taskbar/StashedHandleViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/StashedHandleViewController.java
@@ -57,6 +57,10 @@
     public static final int ALPHA_INDEX_HIDDEN_WHILE_DREAMING = 3;
     private static final int NUM_ALPHA_CHANNELS = 4;
 
+    // Values for long press animations, picked to most closely match navbar spec.
+    private static final float SCALE_TOUCH_ANIMATION_SHRINK = 0.85f;
+    private static final float SCALE_TOUCH_ANIMATION_EXPAND = 1.18f;
+
     /**
      * The SharedPreferences key for whether the stashed handle region is dark.
      */
@@ -324,7 +328,13 @@
 
     @Override
     public void animateNavBarLongPress(boolean isTouchDown, boolean shrink, long durationMs) {
-        // TODO(b/308693847): Animate similarly to NavigationHandle.java (SysUI).
+        float targetScale;
+        if (isTouchDown) {
+            targetScale = shrink ? SCALE_TOUCH_ANIMATION_SHRINK : SCALE_TOUCH_ANIMATION_EXPAND;
+        } else {
+            targetScale = 1f;
+        }
+        mStashedHandleView.animateScale(targetScale, durationMs);
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index e5396ee..b24be54 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -1487,7 +1487,8 @@
             ((LauncherTaskbarUIController) uiController).addLauncherVisibilityChangedAnimation(
                     fullAnimation, duration);
         }
-        mControllers.taskbarStashController.addUnstashToHotseatAnimation(fullAnimation, duration);
+        mControllers.taskbarStashController.addUnstashToHotseatAnimationFromSuw(fullAnimation,
+                duration);
 
         View allAppsButton = mControllers.taskbarViewController.getAllAppsButtonView();
         if (allAppsButton != null && !FeatureFlags.ENABLE_ALL_APPS_BUTTON_IN_HOTSEAT.get()) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index 182ff7e..4da7762 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -355,7 +355,6 @@
         boolean hideTaskbar = isVisible || !mActivity.isUserSetupComplete();
         updateStateForFlag(FLAG_IN_SETUP, hideTaskbar);
         updateStateForFlag(FLAG_STASHED_IN_APP_SETUP, hideTaskbar);
-        updateStateForFlag(FLAG_STASHED_SMALL_SCREEN, mActivity.isPhoneGestureNavMode());
         applyState(hideTaskbar ? 0 : getStashDuration());
     }
 
@@ -547,11 +546,12 @@
      *                            sub-animations are properly coordinated. This duration should not
      *                            actually be used since this animation tracks a swipe progress.
      */
-    protected void addUnstashToHotseatAnimation(AnimatorSet animation, int placeholderDuration) {
+    protected void addUnstashToHotseatAnimationFromSuw(AnimatorSet animation,
+            int placeholderDuration) {
         // Defer any UI updates now to avoid the UI becoming stale when the animation plays.
         mControllers.taskbarViewController.setDeferUpdatesForSUW(true);
         createAnimToIsStashed(
-                /* isStashed= */ false,
+                /* isStashed= */ mActivity.isPhoneMode(),
                 placeholderDuration,
                 TRANSITION_UNSTASH_SUW_MANUAL,
                 /* jankTag= */ "SUW_MANUAL");
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 2ce40b4..d6e4dbd 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -135,6 +135,7 @@
 import com.android.quickstep.util.SwipePipToHomeAnimator;
 import com.android.quickstep.util.TaskViewSimulator;
 import com.android.quickstep.util.TransformParams;
+import com.android.quickstep.views.DesktopTaskView;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.RecentsViewContainer;
 import com.android.quickstep.views.TaskView;
@@ -1258,8 +1259,8 @@
                 ? mRecentsView.getNextPageTaskView() : null;
         TaskView currentPageTaskView = mRecentsView != null
                 ? mRecentsView.getCurrentPageTaskView() : null;
-        if (((nextPageTaskView != null && nextPageTaskView.isDesktopTask())
-                || (currentPageTaskView != null && currentPageTaskView.isDesktopTask()))
+        if ((nextPageTaskView instanceof DesktopTaskView
+                || currentPageTaskView instanceof DesktopTaskView)
                 && endTarget == NEW_TASK) {
             // TODO(b/268075592): add support for quickswitch to/from desktop
             return LAST_TASK;
@@ -1422,7 +1423,7 @@
             setClampScrollOffset(false);
         };
         if (mRecentsView != null && (mRecentsView.getCurrentPageTaskView() != null
-                && !mRecentsView.getCurrentPageTaskView().isDesktopTask())) {
+                && !(mRecentsView.getCurrentPageTaskView() instanceof DesktopTaskView))) {
             ActiveGestureLog.INSTANCE.trackEvent(ActiveGestureErrorDetector.GestureEvent
                     .SET_ON_PAGE_TRANSITION_END_CALLBACK);
             // TODO(b/268075592): add support for quickswitch to/from desktop
@@ -2250,10 +2251,8 @@
                     mRecentsAnimationController, mRecentsAnimationTargets);
         });
 
-        if ((mRecentsView.getNextPageTaskView() != null
-                && mRecentsView.getNextPageTaskView().isDesktopTask())
-                || (mRecentsView.getCurrentPageTaskView() != null
-                && mRecentsView.getCurrentPageTaskView().isDesktopTask())) {
+        if (mRecentsView.getNextPageTaskView() instanceof DesktopTaskView
+                || mRecentsView.getCurrentPageTaskView() instanceof DesktopTaskView) {
             // TODO(b/268075592): add support for quickswitch to/from desktop
             mRecentsViewScrollLinked = false;
             return;
diff --git a/quickstep/src/com/android/quickstep/util/AssistStateManager.java b/quickstep/src/com/android/quickstep/util/AssistStateManager.java
index 42db65f..7acb28d 100644
--- a/quickstep/src/com/android/quickstep/util/AssistStateManager.java
+++ b/quickstep/src/com/android/quickstep/util/AssistStateManager.java
@@ -33,8 +33,8 @@
 
     public AssistStateManager() {}
 
-    /** Whether search is available. */
-    public boolean isSearchAvailable() {
+    /** Return {@code true} if the Settings toggle is enabled. */
+    public boolean isSettingsAllEntrypointsEnabled() {
         return false;
     }
 
@@ -43,14 +43,9 @@
         return false;
     }
 
-    /** Whether CsHelper CtS invocation path is available. */
-    public Optional<Boolean> isCsHelperAvailable() {
-        return Optional.empty();
-    }
-
-    /** Whether VIS CtS invocation path is available. */
-    public Optional<Boolean> isVisAvailable() {
-        return Optional.empty();
+    /** Whether ContextualSearchService invocation path is available. */
+    public boolean isContextualSearchServiceAvailable() {
+        return false;
     }
 
     /** Get the Launcher overridden long press nav handle duration to trigger Assistant. */
@@ -90,11 +85,6 @@
         return Optional.empty();
     }
 
-    /** Return {@code true} if the Settings toggle is enabled. */
-    public boolean isSettingsAllEntrypointsEnabled() {
-        return false;
-    }
-
     /** Dump states. */
     public void dump(String prefix, PrintWriter writer) {}
 
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskView.java b/quickstep/src/com/android/quickstep/views/DesktopTaskView.java
index 4915b62..81c889c 100644
--- a/quickstep/src/com/android/quickstep/views/DesktopTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskView.java
@@ -327,11 +327,6 @@
     }
 
     @Override
-    public boolean isDesktopTask() {
-        return true;
-    }
-
-    @Override
     void refreshThumbnails(@Nullable HashMap<Integer, ThumbnailData> thumbnailDatas) {
         // Sets new thumbnails based on the incoming data and refreshes the rest.
         // Create a copy of the thumbnail map, so we can track thumbnails that need refreshing.
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index f789686..93dd44f 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -1281,11 +1281,6 @@
         return deviceProfile.isTablet && !isFocusedTask();
     }
 
-    /** Whether this task view represents the desktop */
-    public boolean isDesktopTask() {
-        return false;
-    }
-
     /**
      * Called to animate a smooth transition when going directly from an app into Overview (and
      * vice versa). Icons fade in, and DWB banners slide in with a "shift up" animation.
diff --git a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
index efe773b..c64ac23 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
@@ -27,6 +27,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
 
+import com.android.launcher3.allapps.PrivateProfileManager;
 import com.android.launcher3.tapl.HomeAllApps;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.tapl.PrivateSpaceContainer;
@@ -164,6 +165,59 @@
         }
     }
 
+    @Test
+    @ScreenRecordRule.ScreenRecord // b/334946529
+    public void testPrivateSpaceLockingBehaviour() throws IOException {
+        // Scroll to the bottom of All Apps
+        executeOnLauncher(launcher -> launcher.getAppsView().resetAndScrollToPrivateSpaceHeader());
+        HomeAllApps homeAllApps = mLauncher.getAllApps();
+
+        // Disable Private Space
+        togglePrivateSpace(PrivateProfileManager.STATE_DISABLED, homeAllApps);
+
+        homeAllApps.freeze();
+        try {
+            // Verify Locked View elements are present.
+            homeAllApps.getPrivateSpaceLockedView();
+        } finally {
+            // UnFreeze
+            homeAllApps.unfreeze();
+        }
+
+        // Enable Private Space
+        togglePrivateSpace(PrivateProfileManager.STATE_ENABLED, homeAllApps);
+
+        homeAllApps.freeze();
+        try {
+            // Verify UnLocked View elements are present.
+            homeAllApps.getPrivateSpaceUnlockedView();
+        } finally {
+            // UnFreeze
+            homeAllApps.unfreeze();
+        }
+    }
+
+    private void togglePrivateSpace(int state, HomeAllApps homeAllApps) {
+        homeAllApps.freeze();
+        try {
+            // Try Toggling Private Space
+            homeAllApps.togglePrivateSpace();
+        }  finally {
+            // UnFreeze
+            homeAllApps.unfreeze();
+        }
+        PrivateProfileManager manager = getFromLauncher(l -> l.getAppsView()
+                .getPrivateProfileManager());
+        waitForLauncherCondition("Private profile toggle to state: " + state + " failed",
+                launcher -> {
+                    manager.reset();
+                    return manager.getCurrentState() == state;
+                },
+                LauncherInstrumentation.WAIT_TIME_MS);
+        // Wait for Launcher UI to be updated with Private Space Items.
+        waitForLauncherUIUpdate();
+    }
+
     private void waitForPrivateSpaceSetup() {
         waitForLauncherCondition("Private Profile not setup",
                 launcher -> launcher.getAppsView().hasPrivateProfile(),
diff --git a/res/layout/widget_cell_content.xml b/res/layout/widget_cell_content.xml
index 462bb52..b6412db 100644
--- a/res/layout/widget_cell_content.xml
+++ b/res/layout/widget_cell_content.xml
@@ -107,6 +107,7 @@
             android:drawablePadding="@dimen/widget_cell_add_button_drawable_padding"
             android:drawableTint="?attr/widgetPickerAddButtonTextColor"
             android:maxLines="1"
+            style="@style/Button.Rounded.Colored"
             android:background="@drawable/widget_cell_add_button_background" />
     </FrameLayout>
 </merge>
diff --git a/src/com/android/launcher3/allapps/BaseAllAppsAdapter.java b/src/com/android/launcher3/allapps/BaseAllAppsAdapter.java
index a21d7be..16630967 100644
--- a/src/com/android/launcher3/allapps/BaseAllAppsAdapter.java
+++ b/src/com/android/launcher3/allapps/BaseAllAppsAdapter.java
@@ -44,7 +44,6 @@
 import com.android.launcher3.R;
 import com.android.launcher3.allapps.search.SearchAdapterProvider;
 import com.android.launcher3.model.data.AppInfo;
-import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.views.ActivityContext;
 
 /**
@@ -273,8 +272,8 @@
                             privateProfileManager.isPrivateSpaceItem(adapterItem);
                     if (icon.getAlpha() == 0 || icon.getAlpha() == 1) {
                         icon.setAlpha(isPrivateSpaceItem
-                                && privateProfileManager.getAnimationScrolling()
-                                && privateProfileManager.getAnimate()
+                                && (privateProfileManager.getAnimationScrolling() ||
+                                    privateProfileManager.getAnimate())
                                 && privateProfileManager.getCurrentState() == STATE_ENABLED
                                 ? 0 : 1);
                     }
diff --git a/src/com/android/launcher3/allapps/PrivateProfileManager.java b/src/com/android/launcher3/allapps/PrivateProfileManager.java
index 551fa94..ae0e80c 100644
--- a/src/com/android/launcher3/allapps/PrivateProfileManager.java
+++ b/src/com/android/launcher3/allapps/PrivateProfileManager.java
@@ -89,7 +89,15 @@
  */
 public class PrivateProfileManager extends UserProfileManager {
     private static final int EXPAND_COLLAPSE_DURATION = 800;
-    private static final int SETTINGS_OPACITY_DURATION = 160;
+    private static final int SETTINGS_OPACITY_DURATION = 400;
+    private static final int TEXT_UNLOCK_OPACITY_DURATION = 300;
+    private static final int TEXT_LOCK_OPACITY_DURATION = 50;
+    private static final int APP_OPACITY_DURATION = 400;
+    private static final int APP_OPACITY_DELAY = 400;
+    private static final int SETTINGS_AND_LOCK_GROUP_TRANSITION_DELAY = 400;
+    private static final int SETTINGS_OPACITY_DELAY = 400;
+    private static final int LOCK_TEXT_OPACITY_DELAY = 500;
+    private static final int NO_DELAY = 0;
     private final ActivityAllAppsContainerView<?> mAllApps;
     private final Predicate<UserHandle> mPrivateProfileMatcher;
     private final int mPsHeaderHeight;
@@ -445,7 +453,6 @@
         if (getCurrentState() == STATE_ENABLED
                 && isPrivateSpaceSettingsAvailable()) {
             settingsButton.setVisibility(VISIBLE);
-            settingsButton.setAlpha(1f);
             settingsButton.setOnClickListener(
                     view -> {
                         logEvents(LAUNCHER_PRIVATE_SPACE_SETTINGS_TAP);
@@ -590,7 +597,9 @@
         List<BaseAllAppsAdapter.AdapterItem> allAppsAdapterItems =
                 mAllApps.getActiveRecyclerView().getApps().getAdapterItems();
         ValueAnimator alphaAnim = ObjectAnimator.ofFloat(from, to);
-        alphaAnim.setDuration(EXPAND_COLLAPSE_DURATION);
+        alphaAnim.setDuration(APP_OPACITY_DURATION)
+                .setStartDelay(isExpanding ? APP_OPACITY_DELAY : NO_DELAY);
+        alphaAnim.setInterpolator(Interpolators.LINEAR);
         alphaAnim.addUpdateListener(new ValueAnimator.AnimatorUpdateListener() {
             @Override
             public void onAnimationUpdate(ValueAnimator valueAnimator) {
@@ -627,20 +636,25 @@
         }
         ViewGroup settingsAndLockGroup = mPSHeader.findViewById(R.id.settingsAndLockGroup);
         ViewGroup lockButton = mPSHeader.findViewById(R.id.ps_lock_unlock_button);
+        TextView lockText = lockButton.findViewById(R.id.lock_text);
         if (settingsAndLockGroup.getLayoutTransition() == null) {
             // Set a new transition if the current ViewGroup does not already contain one as each
             // transition should only happen once when applied.
             enableLayoutTransition(settingsAndLockGroup);
         }
+        settingsAndLockGroup.getLayoutTransition().setStartDelay(
+                LayoutTransition.CHANGING,
+                expand ? SETTINGS_AND_LOCK_GROUP_TRANSITION_DELAY : NO_DELAY);
         PropertySetter headerSetter = new AnimatedPropertySetter();
         ImageButton settingsButton = mPSHeader.findViewById(R.id.ps_settings_button);
         updateSettingsGearAlpha(settingsButton, expand, headerSetter);
+        updateLockTextAlpha(lockText, expand, headerSetter);
         AnimatorSet animatorSet = headerSetter.buildAnim();
         animatorSet.addListener(new AnimatorListenerAdapter() {
             @Override
             public void onAnimationStart(Animator animation) {
                 // Animate the collapsing of the text at the same time while updating lock button.
-                lockButton.findViewById(R.id.lock_text).setVisibility(expand ? VISIBLE : GONE);
+                lockText.setVisibility(expand ? VISIBLE : GONE);
                 setAnimationRunning(true);
             }
         });
@@ -696,6 +710,8 @@
         LayoutTransition settingsAndLockTransition = new LayoutTransition();
         settingsAndLockTransition.enableTransitionType(LayoutTransition.CHANGING);
         settingsAndLockTransition.setDuration(EXPAND_COLLAPSE_DURATION);
+        settingsAndLockTransition.setInterpolator(LayoutTransition.CHANGING,
+                Interpolators.STANDARD);
         settingsAndLockTransition.addTransitionListener(new LayoutTransition.TransitionListener() {
             @Override
             public void startTransition(LayoutTransition transition, ViewGroup viewGroup,
@@ -716,7 +732,15 @@
             PropertySetter setter) {
         float toAlpha = expand ? 1 : 0;
         setter.setFloat(settingsButton, VIEW_ALPHA, toAlpha, Interpolators.LINEAR)
-                .setDuration(SETTINGS_OPACITY_DURATION).setStartDelay(0);
+                .setDuration(SETTINGS_OPACITY_DURATION).setStartDelay(expand ?
+                        SETTINGS_OPACITY_DELAY : NO_DELAY);
+    }
+
+    private void updateLockTextAlpha(TextView textView, boolean expand, PropertySetter setter) {
+        float toAlpha = expand ? 1 : 0;
+        setter.setFloat(textView, VIEW_ALPHA, toAlpha, Interpolators.LINEAR)
+                .setDuration(expand ? TEXT_UNLOCK_OPACITY_DURATION : TEXT_LOCK_OPACITY_DURATION)
+                .setStartDelay(expand ? LOCK_TEXT_OPACITY_DELAY : NO_DELAY);
     }
 
     void expandPrivateSpace() {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
index 79ddadc..5292ee2 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
@@ -64,7 +64,7 @@
 
     // This ratio defines the max percentage of content area that the recommendations can display
     // with respect to the bottom sheet's height.
-    private static final float RECOMMENDATION_SECTION_HEIGHT_RATIO_TWO_PANE = 0.75f;
+    private static final float RECOMMENDATION_SECTION_HEIGHT_RATIO_TWO_PANE = 0.60f;
     private FrameLayout mSuggestedWidgetsContainer;
     private WidgetsListHeader mSuggestedWidgetsHeader;
     private PackageUserKey mSuggestedWidgetsPackageUserKey;
diff --git a/tests/Android.bp b/tests/Android.bp
index 3822ff8..5ec2263 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -23,7 +23,8 @@
     srcs: [
         "src/**/*.java",
         "src/**/*.kt",
-        ":launcher3-robo-src",
+        "multivalentTests/src/**/*.java",
+        "multivalentTests/src/**/*.kt",
     ],
     exclude_srcs: [
         ":launcher-non-quickstep-tests-src",
@@ -37,6 +38,8 @@
     srcs: [
         "multivalentTests/src/**/*.java",
         "multivalentTests/src/**/*.kt",
+        "src_deviceless/**/*.java",
+        "src_deviceless/**/*.kt",
     ],
 }
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/icons/FastBitmapDrawableTest.java b/tests/multivalentTests/src/com/android/launcher3/icons/FastBitmapDrawableTest.java
index 7e9b68d..58dce0b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/icons/FastBitmapDrawableTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/icons/FastBitmapDrawableTest.java
@@ -37,15 +37,12 @@
 import android.view.animation.PathInterpolator;
 
 import androidx.test.annotation.UiThreadTest;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
-import com.android.launcher3.util.rule.RobolectricUiThreadRule;
+import com.android.launcher3.util.LauncherMultivalentJUnit;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
@@ -55,14 +52,11 @@
  * Tests for FastBitmapDrawable.
  */
 @SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(LauncherMultivalentJUnit.class)
 @UiThreadTest
 public class FastBitmapDrawableTest {
     private static final float EPSILON = 0.00001f;
 
-    @Rule
-    public final TestRule roboUiThreadRule = new RobolectricUiThreadRule();
-
     @Spy
     FastBitmapDrawable mFastBitmapDrawable =
             spy(new FastBitmapDrawable(Bitmap.createBitmap(100, 100, Bitmap.Config.ARGB_8888)));
diff --git a/tests/multivalentTests/src/com/android/launcher3/logging/StartupLatencyLoggerTest.kt b/tests/multivalentTests/src/com/android/launcher3/logging/StartupLatencyLoggerTest.kt
index 12f6c8c..713d4d5 100644
--- a/tests/multivalentTests/src/com/android/launcher3/logging/StartupLatencyLoggerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/logging/StartupLatencyLoggerTest.kt
@@ -2,22 +2,18 @@
 
 import androidx.core.util.isEmpty
 import androidx.test.annotation.UiThreadTest
-import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.launcher3.util.rule.RobolectricUiThreadRule
+import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
 /** Unit test for [ColdRebootStartupLatencyLogger]. */
 @SmallTest
-@RunWith(AndroidJUnit4::class)
+@RunWith(LauncherMultivalentJUnit::class)
 class StartupLatencyLoggerTest {
 
-    @get:Rule val roboUiThreadRule = RobolectricUiThreadRule()
-
     private val underTest = ColdRebootStartupLatencyLogger()
 
     @Before
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/EmulatedDeviceAndroidJUnit.kt b/tests/multivalentTests/src/com/android/launcher3/util/LauncherMultivalentJUnit.kt
similarity index 68%
rename from tests/multivalentTests/src/com/android/launcher3/util/EmulatedDeviceAndroidJUnit.kt
rename to tests/multivalentTests/src/com/android/launcher3/util/LauncherMultivalentJUnit.kt
index 694f257..e8560af 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/EmulatedDeviceAndroidJUnit.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/LauncherMultivalentJUnit.kt
@@ -24,23 +24,24 @@
 import org.junit.runners.Suite
 
 /**
- * A custom runner which emulates multiple devices when running in robolectric framework. Runs
- * normally when running on device
+ * A custom runner for multivalent tests with launcher specific features
+ * 1) Adds support for @UiThread annotations in deviceless tests
+ * 2) Allows emulating multiple devices when running in deviceless mode
  */
-class EmulatedDeviceAndroidJUnit(klass: Class<*>?) : Suite(klass, ImmutableList.of()) {
+class LauncherMultivalentJUnit(klass: Class<*>?) : Suite(klass, ImmutableList.of()) {
 
     val runners: List<Runner> =
-        testClass.getAnnotation(Devices::class.java)?.value?.let { devices ->
-            if (devices.isEmpty() || !isRunningInRobolectric) {
+        (testClass.getAnnotation(EmulatedDevices::class.java)?.value ?: emptyArray()).let { devices
+            ->
+            if (!isRunningInRobolectric) {
                 return@let null
             }
             try {
                 (testClass.javaClass.classLoader.loadClass(ROBOLECTRIC_RUNNER) as Class<Runner>)
                     .getConstructor(Class::class.java, String::class.java)
                     .let { ctor ->
-                        devices.map { deviceName ->
-                            ctor.newInstance(testClass.javaClass, deviceName)
-                        }
+                        if (devices.isEmpty()) listOf(ctor.newInstance(testClass.javaClass, null))
+                        else devices.map { ctor.newInstance(testClass.javaClass, it) }
                     }
             } catch (e: Exception) {
                 null
@@ -50,11 +51,13 @@
 
     override fun getChildren() = runners
 
-    @Retention(RUNTIME) @Target(CLASS) annotation class Devices(val value: Array<String>)
+    /**
+     * Annotation to be added to a test so run it on a list of emulated devices for deviceless test
+     */
+    @Retention(RUNTIME) @Target(CLASS) annotation class EmulatedDevices(val value: Array<String>)
 
     companion object {
-        private const val ROBOLECTRIC_RUNNER =
-            "com.android.launcher3.util.RobolectricEmulatedDeviceRunner"
+        private const val ROBOLECTRIC_RUNNER = "com.android.launcher3.util.RobolectricDeviceRunner"
 
         val isRunningInRobolectric: Boolean
             get() =
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt b/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt
deleted file mode 100644
index b65c443..0000000
--- a/tests/multivalentTests/src/com/android/launcher3/util/rule/RobolectricUiThreadRule.kt
+++ /dev/null
@@ -1,75 +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.util.rule
-
-import androidx.test.annotation.UiThreadTest
-import androidx.test.platform.app.InstrumentationRegistry
-import com.android.launcher3.util.EmulatedDeviceAndroidJUnit.Companion.isRunningInRobolectric
-import java.util.concurrent.atomic.AtomicReference
-import org.junit.rules.TestRule
-import org.junit.runner.Description
-import org.junit.runners.model.Statement
-
-/**
- * A test rule to add support for @UiThreadTest annotations when running in robolectric until is it
- * natively supported by the robolectric runner:
- * https://github.com/robolectric/robolectric/issues/9026
- */
-class RobolectricUiThreadRule : TestRule {
-
-    override fun apply(base: Statement, description: Description): Statement =
-        if (!shouldRunOnUiThread(description)) base else UiThreadStatement(base)
-
-    private fun shouldRunOnUiThread(description: Description): Boolean {
-        if (!isRunningInRobolectric) {
-            // If not running in robolectric, let the default runner handle this
-            return false
-        }
-        var clazz = description.testClass
-        try {
-            if (
-                clazz
-                    .getDeclaredMethod(description.methodName)
-                    .getAnnotation(UiThreadTest::class.java) != null
-            ) {
-                return true
-            }
-        } catch (_: Exception) {
-            // Ignore
-        }
-
-        while (!clazz.isAnnotationPresent(UiThreadTest::class.java)) {
-            clazz = clazz.superclass ?: return false
-        }
-        return true
-    }
-
-    private class UiThreadStatement(val base: Statement) : Statement() {
-
-        override fun evaluate() {
-            val exceptionRef = AtomicReference<Throwable>()
-            InstrumentationRegistry.getInstrumentation().runOnMainSync {
-                try {
-                    base.evaluate()
-                } catch (throwable: Throwable) {
-                    exceptionRef.set(throwable)
-                }
-            }
-            exceptionRef.get()?.let { throw it }
-        }
-    }
-}
diff --git a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
index 70c0333..362596c 100644
--- a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
+++ b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
@@ -177,6 +177,8 @@
      */
     @Test
     @PortraitLandscape
+    @ScreenRecordRule.ScreenRecord // b/338869019
+    @TestStabilityRule.Stability(flavors = LOCAL | PLATFORM_POSTSUBMIT) // b/338869019
     public void testAddDeleteShortcutOnHotseat() {
         mLauncher.getWorkspace()
                 .deleteAppIcon(mLauncher.getWorkspace().getHotseatAppIcon(0))
diff --git a/tests/src_deviceless/com/android/launcher3/util/RobolectricDeviceRunner.kt b/tests/src_deviceless/com/android/launcher3/util/RobolectricDeviceRunner.kt
new file mode 100644
index 0000000..dc6d716
--- /dev/null
+++ b/tests/src_deviceless/com/android/launcher3/util/RobolectricDeviceRunner.kt
@@ -0,0 +1,107 @@
+/*
+ * 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 androidx.test.annotation.UiThreadTest
+import androidx.test.platform.app.InstrumentationRegistry
+import java.lang.reflect.Method
+import java.util.concurrent.atomic.AtomicReference
+import org.junit.runners.model.FrameworkMethod
+import org.junit.runners.model.Statement
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.internal.bytecode.Sandbox
+import org.robolectric.util.ReflectionHelpers
+import org.robolectric.util.ReflectionHelpers.ClassParameter
+
+/** Runner which emulates the provided display before running the actual test */
+class RobolectricDeviceRunner(testClass: Class<*>?, private val deviceName: String?) :
+    RobolectricTestRunner(testClass) {
+
+    private val nameSuffix = deviceName?.let { "-$it" } ?: ""
+
+    override fun getName() = super.getName() + nameSuffix
+
+    override fun testName(method: FrameworkMethod) = super.testName(method) + nameSuffix
+
+    @Throws(Throwable::class)
+    override fun beforeTest(sandbox: Sandbox, method: FrameworkMethod, bootstrappedMethod: Method) {
+        super.beforeTest(sandbox, method, bootstrappedMethod)
+
+        deviceName ?: return
+
+        val emulator =
+            try {
+                ReflectionHelpers.loadClass(
+                    bootstrappedMethod.declaringClass.classLoader,
+                    DEVICE_EMULATOR
+                )
+            } catch (e: Exception) {
+                // Ignore, if the device emulator is not present
+                return
+            }
+        ReflectionHelpers.callStaticMethod<Any>(
+            emulator,
+            "updateDevice",
+            ClassParameter.from(String::class.java, deviceName)
+        )
+    }
+
+    override fun getHelperTestRunner(clazz: Class<*>) = MyHelperTestRunner(clazz)
+
+    class MyHelperTestRunner(clazz: Class<*>) : HelperTestRunner(clazz) {
+
+        override fun methodBlock(method: FrameworkMethod): Statement =
+            // this needs to be run in the test classLoader
+            ReflectionHelpers.callStaticMethod(
+                method.declaringClass.classLoader,
+                RobolectricDeviceRunner::class.qualifiedName,
+                "wrapUiThreadMethod",
+                ClassParameter.from(FrameworkMethod::class.java, method),
+                ClassParameter.from(Statement::class.java, super.methodBlock(method))
+            )
+    }
+
+    private class UiThreadStatement(val base: Statement) : Statement() {
+
+        override fun evaluate() {
+            val exceptionRef = AtomicReference<Throwable>()
+            InstrumentationRegistry.getInstrumentation().runOnMainSync {
+                try {
+                    base.evaluate()
+                } catch (throwable: Throwable) {
+                    exceptionRef.set(throwable)
+                }
+            }
+            exceptionRef.get()?.let { throw it }
+        }
+    }
+
+    companion object {
+
+        private const val DEVICE_EMULATOR = "com.android.launcher3.util.RoboDeviceEmulator"
+
+        @JvmStatic
+        fun wrapUiThreadMethod(method: FrameworkMethod, base: Statement): Statement =
+            if (
+                method.method.isAnnotationPresent(UiThreadTest::class.java) ||
+                    method.declaringClass.isAnnotationPresent(UiThreadTest::class.java)
+            ) {
+                UiThreadStatement(base)
+            } else {
+                base
+            }
+    }
+}
diff --git a/tests/tapl/com/android/launcher3/tapl/AllApps.java b/tests/tapl/com/android/launcher3/tapl/AllApps.java
index f3ec852..71cd19d 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllApps.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllApps.java
@@ -58,6 +58,7 @@
     private static final String FAST_SCROLLER_RES_ID = "fast_scroller";
     private static final Pattern EVENT_ALT_ESC_UP = Pattern.compile(
             "Key event: KeyEvent.*?action=ACTION_UP.*?keyCode=KEYCODE_ESCAPE.*?metaState=0");
+    private static final String UNLOCK_BUTTON_VIEW_RES_ID = "ps_lock_unlock_button";
 
     private final int mHeight;
     private final int mIconHeight;
@@ -433,7 +434,28 @@
     public PrivateSpaceContainer getPrivateSpaceUnlockedView() {
         final UiObject2 allAppsContainer = verifyActiveContainer();
         final UiObject2 appListRecycler = getAppListRecycler(allAppsContainer);
-        return new PrivateSpaceContainer(mLauncher, appListRecycler, this);
+        return new PrivateSpaceContainer(mLauncher, appListRecycler, this, true);
+    }
+
+    /** Returns PrivateSpaceContainer in locked state, if present in view. */
+    @NonNull
+    public PrivateSpaceContainer getPrivateSpaceLockedView() {
+        final UiObject2 allAppsContainer = verifyActiveContainer();
+        final UiObject2 appListRecycler = getAppListRecycler(allAppsContainer);
+        return new PrivateSpaceContainer(mLauncher, appListRecycler, this, false);
+    }
+
+    /**
+     * Toggles Lock/Unlock of Private Space, changing the All Apps Ui.
+     */
+    public void togglePrivateSpace() {
+        final UiObject2 allAppsContainer = verifyActiveContainer();
+        final UiObject2 appListRecycler = getAppListRecycler(allAppsContainer);
+        UiObject2 unLockButtonView = mLauncher.waitForObjectInContainer(appListRecycler,
+                UNLOCK_BUTTON_VIEW_RES_ID);
+        mLauncher.waitForObjectEnabled(unLockButtonView, "Private Space Unlock Button");
+        mLauncher.assertTrue("PS Unlock Button is non-clickable", unLockButtonView.isClickable());
+        unLockButtonView.click();
     }
 
     protected abstract void verifyVisibleContainerOnDismiss();
diff --git a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java
index 0e65ffb..c21bae3 100644
--- a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java
+++ b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java
@@ -31,16 +31,24 @@
     private final LauncherInstrumentation mLauncher;
     private final UiObject2 mAppListRecycler;
     private final AllApps mAppList;
+    private final boolean mPrivateSpaceEnabled;
 
     PrivateSpaceContainer(LauncherInstrumentation launcherInstrumentation,
-            UiObject2 appListRecycler, AllApps appList) {
+            UiObject2 appListRecycler, AllApps appList, boolean privateSpaceEnabled) {
         mLauncher = launcherInstrumentation;
         mAppListRecycler = appListRecycler;
         mAppList = appList;
+        mPrivateSpaceEnabled = privateSpaceEnabled;
 
-        verifyHeaderIsPresent();
-        verifyInstallAppButtonIsPresent();
-        verifyDividerIsPresent();
+        if (mPrivateSpaceEnabled) {
+            verifyHeaderIsPresent();
+            verifyInstallAppButtonIsPresent();
+            verifyDividerIsPresent();
+        } else {
+            verifyHeaderIsPresent();
+            verifyInstallAppButtonIsNotPresent();
+            verifyDividerIsNotPresent();
+        }
     }
 
     // Assert PS Header is in view.
@@ -48,7 +56,7 @@
     private void verifyHeaderIsPresent() {
         final UiObject2 psHeader = mLauncher.waitForObjectInContainer(mAppListRecycler,
                 PS_HEADER_RES_ID);
-        new PrivateSpaceHeader(mLauncher, psHeader, true);
+        new PrivateSpaceHeader(mLauncher, psHeader, mPrivateSpaceEnabled);
     }
 
 
@@ -57,11 +65,21 @@
         mAppList.getAppIcon(INSTALL_APP_TITLE);
     }
 
+    // Assert Install App Item is not present in view.
+    private void verifyInstallAppButtonIsNotPresent() {
+        mLauncher.waitUntilLauncherObjectGone(DIVIDER_RES_ID);
+    }
+
     // Assert Sys App Divider is present in view.
     private void verifyDividerIsPresent() {
         mLauncher.waitForObjectInContainer(mAppListRecycler, DIVIDER_RES_ID);
     }
 
+    // Assert Sys App Divider is not present in view.
+    private void verifyDividerIsNotPresent() {
+        mLauncher.waitUntilLauncherObjectGone(DIVIDER_RES_ID);
+    }
+
     /**
      * Verifies that a user installed app is present above the divider.
      */
@@ -73,4 +91,4 @@
         mLauncher.assertTrue("Installed App: " + appName + " is not above the divider",
                 iconCenter.y < dividerCenter.y);
     }
-}
+}
\ No newline at end of file
diff --git a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceHeader.java b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceHeader.java
index 3a7f038..cc64aae 100644
--- a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceHeader.java
+++ b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceHeader.java
@@ -45,7 +45,7 @@
         if (mPrivateSpaceEnabled) {
             verifyUnlockedState();
         } else {
-            mLauncher.fail("Private Space found in non enabled state");
+            verifyLockedState();
         }
     }
 
@@ -74,4 +74,23 @@
         mLauncher.assertEquals("PS lock text is incorrect", "Lock", lockText.getText());
 
     }
+
+    /** Verify Locked State elements in Private Space Header */
+    private void verifyLockedState() {
+        UiObject2 headerText = mLauncher.waitForObjectInContainer(mPrivateSpaceHeader,
+                PS_HEADER_TEXT_RES_ID);
+        mLauncher.assertEquals("PS Header Text is incorrect ",
+                "Private", headerText.getText());
+
+        UiObject2 unLockButtonView = mLauncher.waitForObjectInContainer(mPrivateSpaceHeader,
+                UNLOCK_BUTTON_VIEW_RES_ID);
+        mLauncher.waitForObjectEnabled(unLockButtonView, "Private Space Unlock Button");
+        mLauncher.assertTrue("PS Unlock Button is non-clickable", unLockButtonView.isClickable());
+
+        mLauncher.waitForObjectInContainer(mPrivateSpaceHeader,
+                LOCK_ICON_RES_ID);
+
+        mLauncher.waitUntilLauncherObjectGone(SETTINGS_BUTTON_RES_ID);
+        mLauncher.waitUntilLauncherObjectGone(LOCK_TEXT_RES_ID);
+    }
 }