Merge "Fix NPE from NavbarButtonsViewController on FEATURE_PC devices." into main
diff --git a/Android.bp b/Android.bp
index b68e1e5..4dddbf6 100644
--- a/Android.bp
+++ b/Android.bp
@@ -136,6 +136,24 @@
     min_sdk_version: min_launcher3_sdk_version,
 }
 
+aconfig_declarations {
+    name: "launcher_flags",
+    package: "com.google.android.platform.launcher.aconfig.flags",
+    srcs: ["launcher.aconfig"],
+}
+
+java_aconfig_library {
+    name: "launcher_flags_lib",
+    aconfig_declarations: "launcher_flags",
+}
+
+java_aconfig_library {
+    name: "launcher_flags_lib_test",
+    aconfig_declarations: "launcher_flags",
+    test: true
+}
+
+
 // Library with all the dependencies for building Launcher3
 android_library {
     name: "Launcher3ResLib",
@@ -167,8 +185,8 @@
 //
 // Build rule for Launcher3 dependencies lib.
 //
-android_library {
-    name: "Launcher3CommonDepsLib",
+java_defaults {
+    name: "Launcher3CommonDepsDefault",
     srcs: ["src_build_config/**/*.java"],
     static_libs: [
         "Launcher3ResLib",
@@ -184,13 +202,35 @@
 }
 
 //
+// Build rule for Launcher3 dependencies lib.
+//
+android_library {
+    name: "Launcher3CommonDepsLib",
+    defaults: ["Launcher3CommonDepsDefault"],
+    static_libs: [
+        "launcher_flags_lib",
+    ],
+}
+
+//
+// Build rule for Launcher3 dependencies lib for test and debug.
+//
+android_library {
+    name: "Launcher3CommonDepsLibDebug",
+    defaults: ["Launcher3CommonDepsDefault"],
+    static_libs: [
+        "launcher_flags_lib_test",
+    ],
+}
+
+//
 // Build rule for Launcher3 app.
 //
 android_app {
     name: "Launcher3",
 
     static_libs: [
-        "Launcher3CommonDepsLib",
+        "Launcher3CommonDepsLibDebug",
     ],
     srcs: [
         ":launcher-src",
diff --git a/build.gradle b/build.gradle
index 090bafe..f4d7261 100644
--- a/build.gradle
+++ b/build.gradle
@@ -157,6 +157,8 @@
     withoutQuickstepImplementation fileTree(dir: "${FRAMEWORK_PREBUILTS_DIR}/libs", include: 'plugin_core.jar')
 
     testImplementation 'junit:junit:4.12'
+    testImplementation libs.mockitoInlineExtended
+    androidTestImplementation libs.mockitoInlineExtended
     androidTestImplementation "org.mockito:mockito-core:1.9.5"
     androidTestImplementation 'com.google.dexmaker:dexmaker:1.2'
     androidTestImplementation 'com.google.dexmaker:dexmaker-mockito:1.2'
diff --git a/launcher.aconfig b/launcher.aconfig
new file mode 100644
index 0000000..cab193c
--- /dev/null
+++ b/launcher.aconfig
@@ -0,0 +1,8 @@
+package: "com.google.android.platform.launcher.aconfig.flags"
+
+flag {
+    name: "enable_expanding_pause_work_button"
+    namespace: "launcher"
+    description: "Expand and collapse pause work button while scrolling."
+    bug: "270390779"
+}
diff --git a/quickstep/res/values-am/strings.xml b/quickstep/res/values-am/strings.xml
index d9e9691..a678b3f 100644
--- a/quickstep/res/values-am/strings.xml
+++ b/quickstep/res/values-am/strings.xml
@@ -92,7 +92,7 @@
     <string name="default_device_name" msgid="6660656727127422487">"መሣሪያ"</string>
     <string name="allset_navigation_settings" msgid="4713404605961476027"><annotation id="link">"የስርዓት አሰሳ ቅንብሮች"</annotation></string>
     <string name="action_share" msgid="2648470652637092375">"አጋራ"</string>
-    <string name="action_screenshot" msgid="8171125848358142917">"ቅጽበታዊ ገፅ እይታ"</string>
+    <string name="action_screenshot" msgid="8171125848358142917">"ቅጽበታዊ ገፅ ዕይታ"</string>
     <string name="action_split" msgid="2098009717623550676">"ክፈል"</string>
     <string name="toast_split_select_app" msgid="8464310533320556058">"የተከፈለ ማያ ገጽን ለመጠቀም ሌላ መተግበሪያ መታ ያድርጉ"</string>
     <string name="toast_split_select_cont_desc" msgid="2119685056059607602">"ከተከፈለ ማያ ገፅ ምርጫ ይውጡ"</string>
diff --git a/quickstep/res/values-iw/strings.xml b/quickstep/res/values-iw/strings.xml
index 2ec3cbf..e07d338 100644
--- a/quickstep/res/values-iw/strings.xml
+++ b/quickstep/res/values-iw/strings.xml
@@ -83,7 +83,7 @@
     <string name="gesture_tutorial_action_button_label" msgid="6249846312991332122">"סיום"</string>
     <string name="gesture_tutorial_action_button_label_settings" msgid="2923621047916486604">"הגדרות"</string>
     <string name="gesture_tutorial_try_again" msgid="65962545858556697">"ניסיון חוזר"</string>
-    <string name="gesture_tutorial_nice" msgid="2936275692616928280">"איזה יופי!"</string>
+    <string name="gesture_tutorial_nice" msgid="2936275692616928280">"יפה!"</string>
     <string name="gesture_tutorial_step" msgid="1279786122817620968">"מדריך <xliff:g id="CURRENT">%1$d</xliff:g>/<xliff:g id="TOTAL">%2$d</xliff:g>"</string>
     <string name="allset_title" msgid="5021126669778966707">"הכול מוכן!"</string>
     <string name="allset_hint" msgid="459504134589971527">"כדי לחזור לדף הבית, מחליקים כלפי מעלה"</string>
diff --git a/quickstep/res/values/override.xml b/quickstep/res/values/override.xml
index 67be0dd..860abc1 100644
--- a/quickstep/res/values/override.xml
+++ b/quickstep/res/values/override.xml
@@ -27,6 +27,8 @@
 
   <string name="nav_handle_long_press_handler_class" translatable="false"></string>
 
+  <string name="assist_utils_class" translatable="false"></string>
+
   <string name="secondary_display_predictions_class" translatable="false">com.android.launcher3.secondarydisplay.SecondaryDisplayPredictionsImpl</string>
 
   <string name="taskbar_model_callbacks_factory_class" translatable="false">com.android.launcher3.taskbar.TaskbarModelCallbacksFactory</string>
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
index 0f8de34..fe8400f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
@@ -51,6 +51,7 @@
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TaskUtils;
 import com.android.quickstep.TouchInteractionService;
+import com.android.quickstep.util.AssistUtils;
 import com.android.quickstep.views.DesktopTaskView;
 
 import java.io.PrintWriter;
@@ -158,7 +159,7 @@
         switch (buttonType) {
             case BUTTON_HOME:
                 logEvent(LAUNCHER_TASKBAR_HOME_BUTTON_LONGPRESS);
-                startAssistant();
+                onLongPressHome();
                 return true;
             case BUTTON_A11Y:
                 logEvent(LAUNCHER_TASKBAR_A11Y_BUTTON_LONGPRESS);
@@ -307,13 +308,17 @@
         }
     }
 
-    private void startAssistant() {
+    private void onLongPressHome() {
         if (mScreenPinned || !mAssistantLongPressEnabled) {
             return;
         }
-        Bundle args = new Bundle();
-        args.putInt(INVOCATION_TYPE_KEY, INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS);
-        mSystemUiProxy.startAssistant(args);
+        // Attempt to start Assist with AssistUtils, otherwise fall back to SysUi's implementation.
+        if (!AssistUtils.newInstance(mService.getApplicationContext()).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);
+        }
     }
 
     private void showQuickSettings() {
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index ffe077b..c482911 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -280,7 +280,7 @@
             // the position of the bubble when the bar is fully expanded
             final float expandedX = i * (mIconSize + mIconSpacing);
             // the position of the bubble when the bar is fully collapsed
-            final float collapsedX = i * mIconOverlapAmount;
+            final float collapsedX = i == 0 ? 0 : mIconOverlapAmount;
 
             if (mIsBarExpanded) {
                 // where the bubble will end up when the animation ends
@@ -292,12 +292,22 @@
                 }
                 // When we're expanded, we're not stacked so we're not behind the stack
                 bv.setBehindStack(false, animate);
+                bv.setAlpha(1);
             } else {
                 final float targetX = currentWidth - collapsedWidth + collapsedX;
                 bv.setTranslationX(widthState * (expandedX - targetX) + targetX);
                 bv.setZ((MAX_BUBBLES * mBubbleElevation) - i);
                 // If we're not the first bubble we're behind the stack
                 bv.setBehindStack(i > 0, animate);
+                // If we're fully collapsed, hide all bubbles except for the first 2. If there are
+                // only 2 bubbles, hide the second bubble as well because it's the overflow.
+                if (widthState == 0) {
+                    if (i > 1) {
+                        bv.setAlpha(0);
+                    } else if (i == 1 && bubbleCount == 2) {
+                        bv.setAlpha(0);
+                    }
+                }
             }
         }
 
@@ -458,7 +468,11 @@
     private float collapsedWidth() {
         final int childCount = getChildCount();
         final int horizontalPadding = getPaddingStart() + getPaddingEnd();
-        return mIconSize + ((childCount - 1) * mIconOverlapAmount) + horizontalPadding;
+        // If there are more than 2 bubbles, the first 2 should be visible when collapsed.
+        // Otherwise just the first bubble should be visible because we don't show the overflow.
+        return childCount > 2
+                ? mIconSize + mIconOverlapAmount + horizontalPadding
+                : mIconSize + horizontalPadding;
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDismissViewExt.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDismissViewExt.kt
index 4b235a9..6c3f0d8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDismissViewExt.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDismissViewExt.kt
@@ -30,6 +30,7 @@
 fun DismissView.setup() {
     setup(
         DismissView.Config(
+            dismissViewResId = R.id.dismiss_view,
             targetSizeResId = R.dimen.bubblebar_dismiss_target_size,
             iconSizeResId = R.dimen.bubblebar_dismiss_target_icon_size,
             bottomMarginResId = R.dimen.bubblebar_dismiss_target_bottom_margin,
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index b444b49..4c701c7 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -132,6 +132,7 @@
 import com.android.launcher3.uioverrides.touchcontrollers.TransposedQuickSwitchTouchController;
 import com.android.launcher3.uioverrides.touchcontrollers.TwoButtonNavbarTouchController;
 import com.android.launcher3.util.ActivityOptionsWrapper;
+import com.android.launcher3.util.BackPressHandler;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.IntSet;
@@ -229,7 +230,8 @@
         mSplitSelectStateController =
                 new SplitSelectStateController(this, mHandler, getStateManager(),
                         getDepthController(), getStatsLogManager(),
-                        SystemUiProxy.INSTANCE.get(this), RecentsModel.INSTANCE.get(this));
+                        SystemUiProxy.INSTANCE.get(this), RecentsModel.INSTANCE.get(this),
+                        () -> onStateBack());
         overviewPanel.init(mActionsView, mSplitSelectStateController);
         mSplitWithKeyboardShortcutController = new SplitWithKeyboardShortcutController(this,
                 mSplitSelectStateController);
@@ -254,6 +256,7 @@
         mEnableWidgetDepth = SystemProperties.getBoolean("ro.launcher.depth.widget", true);
         getWorkspace().addOverlayCallback(progress ->
                 onTaskbarInAppDisplayProgressUpdate(progress, MINUS_ONE_PAGE_PROGRESS_INDEX));
+        addBackAnimationCallback(mSplitSelectStateController.getSplitBackHandler());
     }
 
     @Override
@@ -479,6 +482,7 @@
         mHotseatPredictionController.destroy();
         mSplitWithKeyboardShortcutController.onDestroy();
         if (mViewCapture != null) mViewCapture.close();
+        removeBackAnimationCallback(mSplitSelectStateController.getSplitBackHandler());
     }
 
     @Override
@@ -663,6 +667,10 @@
         anim.buildAnim().start();
     }
 
+    @Override
+    protected boolean isSplitSelectionEnabled() {
+        return mSplitSelectStateController.isSplitSelectActive();
+    }
 
     @Override
     protected void onResume() {
diff --git a/quickstep/src/com/android/quickstep/RecentsActivity.java b/quickstep/src/com/android/quickstep/RecentsActivity.java
index 72439de..d66421f 100644
--- a/quickstep/src/com/android/quickstep/RecentsActivity.java
+++ b/quickstep/src/com/android/quickstep/RecentsActivity.java
@@ -136,7 +136,8 @@
         mSplitSelectStateController =
                 new SplitSelectStateController(this, mHandler, getStateManager(),
                         null /* depthController */, getStatsLogManager(),
-                        SystemUiProxy.INSTANCE.get(this), RecentsModel.INSTANCE.get(this));
+                        SystemUiProxy.INSTANCE.get(this), RecentsModel.INSTANCE.get(this),
+                        null /*activityBackCallback*/);
         mDragLayer.recreateControllers();
         mFallbackRecentsView.init(mActionsView, mSplitSelectStateController);
 
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index ef2d20a..99dd634 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -62,6 +62,7 @@
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Preconditions;
 import com.android.launcher3.util.SplitConfigurationOptions;
+import com.android.quickstep.util.AssistUtils;
 import com.android.systemui.shared.recents.ISystemUiProxy;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.RecentsAnimationControllerCompat;
@@ -250,6 +251,8 @@
         setBackToLauncherCallback(mBackToLauncherCallback, mBackToLauncherRunner);
         setUnfoldAnimationListener(mUnfoldAnimationListener);
         setDesktopTaskListener(mDesktopTaskListener);
+        setAssistantOverridesRequested(
+                AssistUtils.newInstance(mContext).getSysUiAssistOverrideInvocationTypes());
     }
 
     /**
@@ -374,6 +377,17 @@
     }
 
     @Override
+    public void setAssistantOverridesRequested(int[] invocationTypes) {
+        if (mSystemUiProxy != null) {
+            try {
+                mSystemUiProxy.setAssistantOverridesRequested(invocationTypes);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call setAssistantOverridesRequested", e);
+            }
+        }
+    }
+
+    @Override
     public void notifyAccessibilityButtonClicked(int displayId) {
         if (mSystemUiProxy != null) {
             try {
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index 06f1f9a..076f4b1 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -122,6 +122,12 @@
     public void removeListeners() {
     }
 
+    /**
+     * Clears any active state outside of the TaskOverlay lifecycle which might have built
+     * up over time
+     */
+    public void clearAllActiveState() { }
+
     /** Note that these will be shown in order from top to bottom, if available for the task. */
     private static final TaskShortcutFactory[] MENU_OPTIONS = new TaskShortcutFactory[]{
             TaskShortcutFactory.APP_INFO,
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 22aca25..c1680de 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -117,6 +117,7 @@
 import com.android.quickstep.inputconsumers.TrackpadStatusBarInputConsumer;
 import com.android.quickstep.util.ActiveGestureLog;
 import com.android.quickstep.util.ActiveGestureLog.CompoundString;
+import com.android.quickstep.util.AssistUtils;
 import com.android.systemui.shared.recents.IOverviewProxy;
 import com.android.systemui.shared.recents.ISystemUiProxy;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
@@ -280,6 +281,20 @@
             }));
         }
 
+        /**
+         * Sent when the assistant has been invoked with the given type (defined in AssistManager)
+         * and should be shown. This method is used if SystemUiProxy#setAssistantOverridesRequested
+         * was previously called including this invocation type.
+         */
+        @Override
+        public void onAssistantOverrideInvoked(int invocationType) {
+            executeForTouchInteractionService(tis -> {
+                if (!AssistUtils.newInstance(tis).tryStartAssistOverride(invocationType)) {
+                    Log.w(TAG, "Failed to invoke Assist override");
+                }
+            });
+        }
+
         @Override
         public void onNavigationBarSurface(SurfaceControl surface) {
             // TODO: implement
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
index 2816228..4c66504 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
@@ -278,7 +278,8 @@
                     if (!mIsDeferredDownTarget) {
                         // Normal gesture, ensure we pass the drag slop before we start tracking
                         // the gesture
-                        if (Math.abs(displacement) > mTouchSlop) {
+                        if (mGestureState.isTrackpadGesture() || Math.abs(displacement)
+                                > mTouchSlop) {
                             mPassedWindowMoveSlop = true;
                             mStartDisplacement = Math.min(displacement, -mTouchSlop);
                         }
@@ -287,8 +288,8 @@
 
                 float horizontalDist = Math.abs(displacementX);
                 float upDist = -displacement;
-                boolean passedSlop = squaredHypot(displacementX, displacementY)
-                        >= mSquaredTouchSlop;
+                boolean passedSlop = mGestureState.isTrackpadGesture() || squaredHypot(
+                        displacementX, displacementY) >= mSquaredTouchSlop;
 
                 if (!mPassedSlopOnThisGesture && passedSlop) {
                     mPassedSlopOnThisGesture = true;
diff --git a/quickstep/src/com/android/quickstep/util/AssistUtils.java b/quickstep/src/com/android/quickstep/util/AssistUtils.java
new file mode 100644
index 0000000..11b6ea7
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/util/AssistUtils.java
@@ -0,0 +1,45 @@
+/*
+ * 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/SplitSelectStateController.java b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
index 6d5aa16..6865935 100644
--- a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
@@ -42,7 +42,6 @@
 import android.app.ActivityOptions;
 import android.app.ActivityThread;
 import android.app.PendingIntent;
-import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
@@ -77,8 +76,10 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.statehandlers.DepthController;
 import com.android.launcher3.statemanager.StateManager;
+import com.android.launcher3.statemanager.StatefulActivity;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.shared.TestProtocol;
+import com.android.launcher3.util.BackPressHandler;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.SplitConfigurationOptions.StagePosition;
 import com.android.quickstep.OverviewComponentObserver;
@@ -113,9 +114,11 @@
 public class SplitSelectStateController {
     private static final String TAG = "SplitSelectStateCtor";
 
-    private Context mContext;
+    private StatefulActivity mContext;
     private final Handler mHandler;
     private final RecentsModel mRecentTasksModel;
+    @Nullable
+    private final Runnable mActivityBackCallback;
     private final SplitAnimationController mSplitAnimationController;
     private final AppPairsController mAppPairsController;
     private final SplitSelectDataHolder mSplitSelectDataHolder;
@@ -142,9 +145,28 @@
 
     private final List<SplitSelectionListener> mSplitSelectionListeners = new ArrayList<>();
 
-    public SplitSelectStateController(Context context, Handler handler, StateManager stateManager,
-            DepthController depthController, StatsLogManager statsLogManager,
-            SystemUiProxy systemUiProxy, RecentsModel recentsModel) {
+    private final BackPressHandler mSplitBackHandler = new BackPressHandler() {
+        @Override
+        public boolean canHandleBack() {
+            return FeatureFlags.ENABLE_SPLIT_FROM_WORKSPACE_TO_WORKSPACE.get() &&
+                    isSplitSelectActive();
+        }
+
+        @Override
+        public void onBackInvoked() {
+            // When exiting from split selection, leave current context to go to
+            // homescreen as well
+            getSplitAnimationController().playPlaceholderDismissAnim(mContext);
+            if (mActivityBackCallback != null) {
+                mActivityBackCallback.run();
+            }
+        }
+    };
+
+    public SplitSelectStateController(StatefulActivity context, Handler handler,
+            StateManager stateManager, DepthController depthController,
+            StatsLogManager statsLogManager, SystemUiProxy systemUiProxy, RecentsModel recentsModel,
+            Runnable activityBackCallback) {
         mContext = context;
         mHandler = handler;
         mStatsLogManager = statsLogManager;
@@ -152,6 +174,7 @@
         mStateManager = stateManager;
         mDepthController = depthController;
         mRecentTasksModel = recentsModel;
+        mActivityBackCallback = activityBackCallback;
         mSplitAnimationController = new SplitAnimationController(this);
         mAppPairsController = new AppPairsController(context, this, statsLogManager);
         mSplitSelectDataHolder = new SplitSelectDataHolder(mContext);
@@ -716,6 +739,10 @@
         return mAppPairsController;
     }
 
+    public BackPressHandler getSplitBackHandler() {
+        return mSplitBackHandler;
+    }
+
     public void dump(String prefix, PrintWriter writer) {
         if (mSplitSelectDataHolder != null) {
             mSplitSelectDataHolder.dump(prefix, writer);
diff --git a/quickstep/src/com/android/quickstep/views/FloatingTaskView.java b/quickstep/src/com/android/quickstep/views/FloatingTaskView.java
index f250b8c..92ffcd0 100644
--- a/quickstep/src/com/android/quickstep/views/FloatingTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/FloatingTaskView.java
@@ -1,3 +1,18 @@
+/*
+ * 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.views;
 
 import static com.android.app.animation.Interpolators.LINEAR;
@@ -47,8 +62,6 @@
  * {@link #addConfirmAnimation(PendingAnimation, RectF, Rect, boolean, boolean)}
  * giving a starting and ending bounds. Currently this is set to use the split placeholder view,
  * but it could be generified.
- *
- * TODO: Figure out how to copy thumbnail data from existing TaskView to this view.
  */
 public class FloatingTaskView extends FrameLayout {
 
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index cb5b457..be9da34 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1396,6 +1396,7 @@
             // its thumbnail
             mTmpRunningTasks = null;
             mSplitBoundsConfig = null;
+            mTaskOverlayFactory.clearAllActiveState();
         }
         updateLocusId();
     }
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
index 9622619..b3d04c6 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
@@ -26,6 +26,7 @@
 import android.os.Handler;
 import android.view.View;
 
+import androidx.test.platform.app.InstrumentationRegistry;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.launcher3.logging.StatsLogManager;
@@ -70,6 +71,9 @@
         MockitoAnnotations.initMocks(this);
         when(mockService.getDisplayId()).thenReturn(DISPLAY_ID);
         when(mockService.getOverviewCommandHelper()).thenReturn(mockCommandHelper);
+        when(mockService.getApplicationContext())
+                .thenReturn(InstrumentationRegistry.getInstrumentation().getTargetContext()
+                        .getApplicationContext());
         when(mockStatsLogManager.logger()).thenReturn(mockStatsLogger);
         when(mockTaskbarControllers.getTaskbarActivityContext())
                 .thenReturn(mockTaskbarActivityContext);
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsSplitscreen.java b/quickstep/tests/src/com/android/quickstep/TaplTestsSplitscreen.java
index 1aa7ab6..92b598b 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsSplitscreen.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsSplitscreen.java
@@ -25,6 +25,8 @@
 
 import android.content.Intent;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.LargeTest;
 import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.launcher3.config.FeatureFlags;
@@ -36,7 +38,10 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
 
+@LargeTest
+@RunWith(AndroidJUnit4.class)
 public class TaplTestsSplitscreen extends AbstractQuickStepTest {
     private static final String CALCULATOR_APP_NAME = "Calculator";
     private static final String CALCULATOR_APP_PACKAGE =
diff --git a/quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt b/quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
index 69109c2..4e7fcf0 100644
--- a/quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
@@ -31,6 +31,7 @@
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.statehandlers.DepthController
 import com.android.launcher3.statemanager.StateManager
+import com.android.launcher3.statemanager.StatefulActivity
 import com.android.launcher3.util.ComponentKey
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.launcher3.util.withArgCaptor
@@ -58,7 +59,7 @@
     @Mock lateinit var statsLogManager: StatsLogManager
     @Mock lateinit var stateManager: StateManager<LauncherState>
     @Mock lateinit var handler: Handler
-    @Mock lateinit var context: Context
+    @Mock lateinit var context: StatefulActivity<*>
     @Mock lateinit var recentsModel: RecentsModel
     @Mock lateinit var pendingIntent: PendingIntent
 
@@ -81,7 +82,8 @@
                 depthController,
                 statsLogManager,
                 systemUiProxy,
-                recentsModel
+                recentsModel,
+                null /*activityBackCallback*/
             )
     }
 
diff --git a/res/values/id.xml b/res/values/id.xml
index 7b812de..872ae2f 100644
--- a/res/values/id.xml
+++ b/res/values/id.xml
@@ -43,4 +43,6 @@
     <item type="id" name="saved_floating_widget_foreground" />
     <item type="id" name="saved_floating_widget_background" />
 
+    <item type="id" name="dismiss_view" />
+
 </resources>
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index ffb8b82..439cc00 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -199,6 +199,7 @@
 import com.android.launcher3.util.ActivityResultInfo;
 import com.android.launcher3.util.ActivityTracker;
 import com.android.launcher3.util.CannedAnimationCoordinator;
+import com.android.launcher3.util.BackPressHandler;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
@@ -425,6 +426,8 @@
     private final CannedAnimationCoordinator mAnimationCoordinator =
             new CannedAnimationCoordinator(this);
 
+    private final List<BackPressHandler> mBackPressedHandlers = new ArrayList<>();
+
     @Override
     @TargetApi(Build.VERSION_CODES.S)
     protected void onCreate(Bundle savedInstanceState) {
@@ -604,6 +607,7 @@
      *  <li> auto cancel action mode handler
      *  <li> drag handler
      *  <li> view handler
+     *  <li> registered {@link BackPressHandler}
      *  <li> state handler
      * </ol>
      *
@@ -633,7 +637,14 @@
             return topView;
         }
 
-        // #4 state handler
+        // #4 Custom back handlers
+        for (BackPressHandler handler : mBackPressedHandlers) {
+            if (handler.canHandleBack()) {
+                return handler;
+            }
+        }
+
+        // #5 state handler
         return new OnBackAnimationCallback() {
             @Override
             public void onBackInvoked() {
@@ -3281,6 +3292,14 @@
         updateDisallowBack();
     }
 
+    protected void addBackAnimationCallback(BackPressHandler callback) {
+        mBackPressedHandlers.add(callback);
+    }
+
+    protected void removeBackAnimationCallback(BackPressHandler callback) {
+        mBackPressedHandlers.remove(callback);
+    }
+
     private void updateDisallowBack() {
         if (DESKTOP_MODE_1_SUPPORTED || DESKTOP_MODE_2_SUPPORTED) {
             // Do not disable back in launcher when prototype behavior is enabled
@@ -3288,12 +3307,20 @@
         }
         LauncherRootView rv = getRootView();
         if (rv != null) {
+            boolean isSplitSelectionEnabled = isSplitSelectionEnabled();
             boolean disableBack = getStateManager().getState() == NORMAL
-                    && AbstractFloatingView.getTopOpenView(this) == null;
+                    && AbstractFloatingView.getTopOpenView(this) == null
+                    && !isSplitSelectionEnabled;
             rv.setDisallowBackGesture(disableBack);
         }
     }
 
+    /** To be overridden by subclasses */
+    protected boolean isSplitSelectionEnabled() {
+        // Overridden
+        return false;
+    }
+
     @Override
     public void returnToHomescreen() {
         super.returnToHomescreen();
diff --git a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
index 40382b2..542266a 100644
--- a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
@@ -26,6 +26,8 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.ScrollableLayoutManager.PREDICTIVE_BACK_MIN_SCALE;
 
+import static com.google.android.platform.launcher.aconfig.flags.Flags.enableExpandingPauseWorkButton;
+
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ValueAnimator;
@@ -561,7 +563,8 @@
             mAH.get(AdapterHolder.MAIN).setup(mainRecyclerView, mPersonalMatcher);
             mAH.get(AdapterHolder.WORK).setup(workRecyclerView, mWorkManager.getMatcher());
             workRecyclerView.setId(R.id.apps_list_view_work);
-            if (FeatureFlags.ENABLE_EXPANDING_PAUSE_WORK_BUTTON.get()) {
+            if (enableExpandingPauseWorkButton()
+                    || FeatureFlags.ENABLE_EXPANDING_PAUSE_WORK_BUTTON.get()) {
                 mAH.get(AdapterHolder.WORK).mRecyclerView.addOnScrollListener(
                         mWorkManager.newScrollListener());
             }
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 9bc2a0a..21520bf 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -230,6 +230,7 @@
     public static final BooleanFlag ENABLE_HIDE_HEADER = getReleaseFlag(270390930,
             "ENABLE_HIDE_HEADER", ENABLED, "Hide header on keyboard before typing in all apps");
 
+    // Aconfig migration complete for ENABLE_EXPANDING_PAUSE_WORK_BUTTON.
     public static final BooleanFlag ENABLE_EXPANDING_PAUSE_WORK_BUTTON = getDebugFlag(270390779,
             "ENABLE_EXPANDING_PAUSE_WORK_BUTTON", DISABLED,
             "Expand and collapse pause work button while scrolling");
diff --git a/src/com/android/launcher3/util/BackPressHandler.java b/src/com/android/launcher3/util/BackPressHandler.java
new file mode 100644
index 0000000..b63f648
--- /dev/null
+++ b/src/com/android/launcher3/util/BackPressHandler.java
@@ -0,0 +1,30 @@
+/*
+ * 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.launcher3.util;
+
+import android.os.Build;
+import android.window.OnBackAnimationCallback;
+
+import androidx.annotation.RequiresApi;
+
+/**
+ * Extension of {@link OnBackAnimationCallback} that allows a check to determine
+ * if this callback supports handling back or not
+ */
+@RequiresApi(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+public interface BackPressHandler extends OnBackAnimationCallback {
+    boolean canHandleBack();
+}
diff --git a/tests/Android.bp b/tests/Android.bp
index 5a52440..89ebe8c 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -73,6 +73,7 @@
     asset_dirs: ["assets"],
     resource_dirs: ["res"],
     static_libs: [
+        "flag-junit-base",
         "launcher-aosp-tapl",
         "androidx.test.core",
         "androidx.test.runner",
@@ -82,11 +83,12 @@
         "androidx.test.espresso.contrib",
         "androidx.test.espresso.intents",
         "androidx.test.uiautomator_uiautomator",
-        "mockito-target-inline-minus-junit4",
+        "mockito-target-extended-minus-junit4",
         "launcher_log_protos_lite",
         "truth-prebuilt",
         "platform-test-rules",
         "testables",
+        "launcher_flags_lib_test",
     ],
     manifest: "AndroidManifest-common.xml",
     platform_apis: true,
@@ -103,12 +105,16 @@
         ":launcher-tests-src",
         ":launcher-non-quickstep-tests-src",
     ],
-    static_libs: ["Launcher3TestLib"],
+    static_libs: [
+        "Launcher3TestLib",
+        "launcher_flags_lib_test",
+    ],
     libs: [
         "android.test.base",
         "android.test.runner",
         "android.test.mock",
     ],
+    // Libraries used by mockito inline extended
     jni_libs: [
         "libdexmakerjvmtiagent",
         "libstaticjvmtiagent",
diff --git a/tests/src/com/android/launcher3/ui/WorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
index 5b9adcd..61cdd17 100644
--- a/tests/src/com/android/launcher3/ui/WorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -22,10 +22,13 @@
 import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
 import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
 
+import static com.google.android.platform.launcher.aconfig.flags.Flags.FLAG_ENABLE_EXPANDING_PAUSE_WORK_BUTTON;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
 
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.util.Log;
 import android.view.View;
 
@@ -44,6 +47,7 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Ignore;
+import org.junit.Rule;
 import org.junit.Test;
 
 import java.util.Objects;
@@ -52,6 +56,7 @@
 public class WorkProfileTest extends AbstractLauncherUiTest {
 
     private static final int WORK_PAGE = ActivityAllAppsContainerView.AdapterHolder.WORK;
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     private int mProfileUserId;
     private boolean mWorkProfileSetupSuccessful;
@@ -60,6 +65,7 @@
     @Before
     @Override
     public void setUp() throws Exception {
+        mSetFlagsRule.disableFlags(FLAG_ENABLE_EXPANDING_PAUSE_WORK_BUTTON);
         super.setUp();
         String output =
                 mDevice.executeShellCommand(
diff --git a/tests/src/com/android/launcher3/util/rule/StaticMockitoRule.java b/tests/src/com/android/launcher3/util/rule/StaticMockitoRule.java
new file mode 100644
index 0000000..6b91474
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/rule/StaticMockitoRule.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2022 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 static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+
+import com.android.dx.mockito.inline.extended.StaticMockitoSession;
+import com.android.dx.mockito.inline.extended.StaticMockitoSessionBuilder;
+
+import org.junit.rules.MethodRule;
+import org.junit.runners.model.FrameworkMethod;
+import org.junit.runners.model.Statement;
+import org.mockito.junit.MockitoRule;
+import org.mockito.quality.Strictness;
+
+/**
+ * Similar to {@link MockitoRule}, but uses {@link StaticMockitoSession}, which allows mocking
+ * static methods.
+ */
+public class StaticMockitoRule implements MethodRule {
+    private Class<?>[] mClasses;
+
+    public StaticMockitoRule(Class<?>... classes) {
+        mClasses = classes;
+    }
+
+    @Override
+    public Statement apply(Statement base, FrameworkMethod method, Object target) {
+        return new Statement() {
+            public void evaluate() throws Throwable {
+                StaticMockitoSessionBuilder builder =
+                        mockitoSession()
+                                .name(target.getClass().getSimpleName() + "." + method.getName())
+                                .initMocks(target)
+                                .strictness(Strictness.STRICT_STUBS);
+
+                for (Class<?> clazz : mClasses) {
+                    builder.mockStatic(clazz);
+                }
+
+                StaticMockitoSession session = builder.startMocking();
+                Throwable testFailure = evaluateSafely(base);
+                session.finishMocking(testFailure);
+                if (testFailure != null) {
+                    throw testFailure;
+                }
+            }
+
+            private Throwable evaluateSafely(Statement base) {
+                try {
+                    base.evaluate();
+                    return null;
+                } catch (Throwable throwable) {
+                    return throwable;
+                }
+            }
+        };
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/viewcapture_analysis/FlashDetector.java b/tests/src/com/android/launcher3/util/viewcapture_analysis/FlashDetector.java
index 388a59a..6d9198f 100644
--- a/tests/src/com/android/launcher3/util/viewcapture_analysis/FlashDetector.java
+++ b/tests/src/com/android/launcher3/util/viewcapture_analysis/FlashDetector.java
@@ -59,6 +59,11 @@
             DRAG_LAYER + "FloatingTaskView|FloatingTaskThumbnailView:id/thumbnail",
             DRAG_LAYER
                     + "WidgetsFullSheet|SpringRelativeLayout:id/container|WidgetsRecyclerView:id"
+                    + "/primary_widgets_list_view|WidgetsListHeader:id/widgets_list_header",
+            DRAG_LAYER
+                    + "WidgetsTwoPaneSheet|SpringRelativeLayout:id/container|LinearLayout:id"
+                    + "/linear_layout_container|FrameLayout:id/recycler_view_container"
+                    + "|FrameLayout:id/widgets_two_pane_sheet_recyclerview|WidgetsRecyclerView:id"
                     + "/primary_widgets_list_view|WidgetsListHeader:id/widgets_list_header"
     ));
 
diff --git a/tests/tapl/com/android/launcher3/tapl/AllApps.java b/tests/tapl/com/android/launcher3/tapl/AllApps.java
index 23d09d4..fb08ea4 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllApps.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllApps.java
@@ -336,4 +336,11 @@
         final Bundle testInfo = mLauncher.getTestInfo(TestProtocol.REQUEST_APP_LIST_FREEZE_FLAGS);
         return testInfo == null ? 0 : testInfo.getInt(TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
+
+    /**
+     * Return the QSB UI object on the AllApps screen.
+     * @return the QSB UI object.
+     */
+    @NonNull
+    public abstract Qsb getQsb();
 }
\ No newline at end of file
diff --git a/tests/tapl/com/android/launcher3/tapl/AllAppsFromTaskbar.java b/tests/tapl/com/android/launcher3/tapl/AllAppsFromTaskbar.java
index c4744a1..0e0291f 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllAppsFromTaskbar.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllAppsFromTaskbar.java
@@ -62,4 +62,10 @@
         return mLauncher.getTestInfo(TestProtocol.REQUEST_TASKBAR_APPS_LIST_SCROLL_Y)
                 .getInt(TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
+
+    @NonNull
+    @Override
+    public TaskbarAllAppsQsb getQsb() {
+        return new TaskbarAllAppsQsb(mLauncher, verifyActiveContainer());
+    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/AllAppsQsb.java b/tests/tapl/com/android/launcher3/tapl/AllAppsQsb.java
index 0931cd4..1692351 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllAppsQsb.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllAppsQsb.java
@@ -22,16 +22,7 @@
  */
 class AllAppsQsb extends Qsb {
 
-    private final UiObject2 mAllAppsContainer;
-
     AllAppsQsb(LauncherInstrumentation launcher, UiObject2 allAppsContainer) {
-        super(launcher);
-        mAllAppsContainer = allAppsContainer;
-        waitForQsbObject();
-    }
-
-    @Override
-    protected UiObject2 waitForQsbObject() {
-        return mLauncher.waitForObjectInContainer(mAllAppsContainer, "search_container_all_apps");
+        super(launcher, allAppsContainer, "search_container_all_apps");
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/HomeAllApps.java b/tests/tapl/com/android/launcher3/tapl/HomeAllApps.java
index a03472a..33c6334 100644
--- a/tests/tapl/com/android/launcher3/tapl/HomeAllApps.java
+++ b/tests/tapl/com/android/launcher3/tapl/HomeAllApps.java
@@ -117,11 +117,8 @@
         }
     }
 
-    /**
-     * Return the QSB UI object on the AllApps screen.
-     * @return the QSB UI object.
-     */
     @NonNull
+    @Override
     public Qsb getQsb() {
         return new AllAppsQsb(mLauncher, verifyActiveContainer());
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/HomeQsb.java b/tests/tapl/com/android/launcher3/tapl/HomeQsb.java
index 20d09a1..5385c65 100644
--- a/tests/tapl/com/android/launcher3/tapl/HomeQsb.java
+++ b/tests/tapl/com/android/launcher3/tapl/HomeQsb.java
@@ -22,16 +22,7 @@
  */
 class HomeQsb extends Qsb {
 
-    private final UiObject2 mHotSeat;
-
     HomeQsb(LauncherInstrumentation launcher, UiObject2 hotseat) {
-        super(launcher);
-        mHotSeat = hotseat;
-        waitForQsbObject();
-    }
-
-    @Override
-    protected UiObject2 waitForQsbObject() {
-        return mLauncher.waitForObjectInContainer(mHotSeat, "search_container_hotseat");
+        super(launcher, hotseat, "search_container_hotseat");
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/Qsb.java b/tests/tapl/com/android/launcher3/tapl/Qsb.java
index 6bc4f21..7f3f61d 100644
--- a/tests/tapl/com/android/launcher3/tapl/Qsb.java
+++ b/tests/tapl/com/android/launcher3/tapl/Qsb.java
@@ -30,13 +30,21 @@
     private static final String ASSISTANT_APP_PACKAGE = "com.google.android.googlequicksearchbox";
     private static final String ASSISTANT_ICON_RES_ID = "mic_icon";
     protected final LauncherInstrumentation mLauncher;
+    private final UiObject2 mContainer;
+    private final String mQsbResName;
 
-    protected Qsb(LauncherInstrumentation launcher) {
+    protected Qsb(LauncherInstrumentation launcher, UiObject2 container, String qsbResName) {
         mLauncher = launcher;
+        mContainer = container;
+        mQsbResName = qsbResName;
+        waitForQsbObject();
     }
 
     // Waits for the quick search box.
-    protected abstract UiObject2 waitForQsbObject();
+    private UiObject2 waitForQsbObject() {
+        return mLauncher.waitForObjectInContainer(mContainer, mQsbResName);
+    }
+
     /**
      * Launch assistant app by tapping mic icon on qsb.
      */
@@ -79,8 +87,12 @@
             mLauncher.waitForIdle();
             try (LauncherInstrumentation.Closable c2 = mLauncher.addContextLayer(
                     "clicked qsb to open search result page")) {
-                return new SearchResultFromQsb(mLauncher);
+                return createSearchResult();
             }
         }
     }
+
+    protected SearchResultFromQsb createSearchResult() {
+        return new SearchResultFromQsb(mLauncher);
+    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java b/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java
index 80176e9..8c3402f 100644
--- a/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java
+++ b/tests/tapl/com/android/launcher3/tapl/SearchResultFromQsb.java
@@ -32,7 +32,7 @@
 
     // This particular ID change should happen with caution
     private static final String SEARCH_CONTAINER_RES_ID = "search_results_list_view";
-    private final LauncherInstrumentation mLauncher;
+    protected final LauncherInstrumentation mLauncher;
 
     SearchResultFromQsb(LauncherInstrumentation launcher) {
         mLauncher = launcher;
@@ -49,8 +49,12 @@
     }
 
     /** Find the app from search results with app name. */
-    public Launchable findAppIcon(String appName) {
+    public AppIcon findAppIcon(String appName) {
         UiObject2 icon = mLauncher.waitForLauncherObject(By.clazz(TextView.class).text(appName));
+        return createAppIcon(icon);
+    }
+
+    protected AppIcon createAppIcon(UiObject2 icon) {
         return new AllAppsAppIcon(mLauncher, icon);
     }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/SearchResultFromTaskbarQsb.java b/tests/tapl/com/android/launcher3/tapl/SearchResultFromTaskbarQsb.java
new file mode 100644
index 0000000..c267c9e
--- /dev/null
+++ b/tests/tapl/com/android/launcher3/tapl/SearchResultFromTaskbarQsb.java
@@ -0,0 +1,38 @@
+/*
+ * 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.launcher3.tapl;
+
+import androidx.test.uiautomator.UiObject2;
+
+/**
+ * Operations on search result page opened from Taskbar qsb.
+ */
+public class SearchResultFromTaskbarQsb extends SearchResultFromQsb {
+
+    SearchResultFromTaskbarQsb(LauncherInstrumentation launcher) {
+        super(launcher);
+    }
+
+    @Override
+    public TaskbarAppIcon findAppIcon(String appName) {
+        return (TaskbarAppIcon) super.findAppIcon(appName);
+    }
+
+    @Override
+    protected TaskbarAppIcon createAppIcon(UiObject2 icon) {
+        return new TaskbarAppIcon(mLauncher, icon);
+    }
+}
diff --git a/tests/tapl/com/android/launcher3/tapl/TaskbarAllAppsQsb.java b/tests/tapl/com/android/launcher3/tapl/TaskbarAllAppsQsb.java
new file mode 100644
index 0000000..7cecd3e
--- /dev/null
+++ b/tests/tapl/com/android/launcher3/tapl/TaskbarAllAppsQsb.java
@@ -0,0 +1,38 @@
+/*
+ * 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.launcher3.tapl;
+
+import androidx.test.uiautomator.UiObject2;
+
+/**
+ * Operations on Taskbar AllApp screen qsb.
+ */
+public class TaskbarAllAppsQsb extends Qsb {
+
+    TaskbarAllAppsQsb(LauncherInstrumentation launcher, UiObject2 allAppsContainer) {
+        super(launcher, allAppsContainer, "search_container_all_apps");
+    }
+
+    @Override
+    public SearchResultFromTaskbarQsb showSearchResult() {
+        return (SearchResultFromTaskbarQsb) super.showSearchResult();
+    }
+
+    @Override
+    protected SearchResultFromTaskbarQsb createSearchResult() {
+        return new SearchResultFromTaskbarQsb(mLauncher);
+    }
+}