diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index 878aa6e..1856b39 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -526,4 +526,11 @@
     namespace: "launcher"
     description: "Enable Taskbar LayoutTransition for Recent Apps"
     bug: "343521765"
+}
+
+flag {
+    name: "enable_pinning_app_with_context_menu"
+    namespace: "launcher"
+    description: "Add options to pin/unpin to taskbar to app context menus."
+    bug: "375648361"
 }
\ No newline at end of file
diff --git a/quickstep/AndroidManifest.xml b/quickstep/AndroidManifest.xml
index 57bfb4a..8c39585 100644
--- a/quickstep/AndroidManifest.xml
+++ b/quickstep/AndroidManifest.xml
@@ -152,7 +152,7 @@
             android:showOnLockScreen="true"
             android:launchMode="singleTop"
             android:exported="true"
-            android:permission="android.permission.START_WIDGET_PICKER_ACTIVITY">
+            android:permission="${applicationId}.permission.START_WIDGET_PICKER_ACTIVITY">
             <intent-filter>
                 <action android:name="android.intent.action.PICK" />
                 <category android:name="android.intent.category.DEFAULT" />
diff --git a/quickstep/res/layout/keyboard_quick_switch_view.xml b/quickstep/res/layout/keyboard_quick_switch_view.xml
index 2420a46..4118500 100644
--- a/quickstep/res/layout/keyboard_quick_switch_view.xml
+++ b/quickstep/res/layout/keyboard_quick_switch_view.xml
@@ -22,6 +22,7 @@
     android:layout_height="wrap_content"
     android:layout_marginTop="@dimen/keyboard_quick_switch_margin_top"
     android:layout_marginHorizontal="@dimen/keyboard_quick_switch_margin_ends"
+    android:layout_gravity="center_horizontal"
     android:background="@drawable/keyboard_quick_switch_view_background"
     android:clipToOutline="true"
     android:alpha="0"
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index 6367a01..70efe14 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -118,6 +118,7 @@
 
     <!-- Launcher app transition -->
     <dimen name="closing_window_trans_y">115dp</dimen>
+    <dimen name="closing_freeform_window_trans_y">36dp</dimen>
 
     <dimen name="quick_switch_scaling_scroll_threshold">100dp</dimen>
 
diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index 18337d3..3bf5f3b 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -108,6 +108,7 @@
 import android.view.animation.AnimationUtils;
 import android.view.animation.Interpolator;
 import android.view.animation.PathInterpolator;
+import android.window.DesktopModeFlags;
 import android.window.RemoteTransition;
 import android.window.TransitionFilter;
 import android.window.WindowAnimationState;
@@ -166,11 +167,13 @@
 import com.android.systemui.shared.system.BlurUtils;
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
 import com.android.systemui.shared.system.QuickStepContract;
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.startingsurface.IStartingWindowListener;
 
 import java.io.PrintWriter;
 import java.lang.ref.WeakReference;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -233,6 +236,7 @@
     protected final Handler mHandler;
 
     private final float mClosingWindowTransY;
+    private final float mClosingFreeformWindowTransY;
     private final float mMaxShadowRadius;
 
     private final StartingWindowListener mStartingWindowListener =
@@ -290,6 +294,8 @@
 
         Resources res = mLauncher.getResources();
         mClosingWindowTransY = res.getDimensionPixelSize(R.dimen.closing_window_trans_y);
+        mClosingFreeformWindowTransY =
+                res.getDimensionPixelSize(R.dimen.closing_freeform_window_trans_y);
         mMaxShadowRadius = res.getDimensionPixelSize(R.dimen.max_shadow_radius);
 
         mLauncher.addOnDeviceProfileChangeListener(this);
@@ -1480,10 +1486,16 @@
                 ? 0 : getWindowCornerRadius(mLauncher);
         float startShadowRadius = areAllTargetsTranslucent(appTargets) ? 0 : mMaxShadowRadius;
         closingAnimator.setDuration(duration);
+        boolean isFreeform = isFreeformAnimation(appTargets);
+        float translateY = isFreeform ? mClosingFreeformWindowTransY : mClosingWindowTransY;
+        float endScale = isFreeform ? 0.95f : 1f;
+        Interpolator alphaInterpolator = isFreeform
+                ? clampToDuration(LINEAR, 0, 100, duration)
+                : clampToDuration(LINEAR, 25, 125, duration);
         closingAnimator.addUpdateListener(new MultiValueUpdateListener() {
-            FloatProp mDy = new FloatProp(0, mClosingWindowTransY, DECELERATE_1_7);
-            FloatProp mScale = new FloatProp(1f, 1f, DECELERATE_1_7);
-            FloatProp mAlpha = new FloatProp(1f, 0f, clampToDuration(LINEAR, 25, 125, duration));
+            FloatProp mDy = new FloatProp(0, translateY, DECELERATE_1_7);
+            FloatProp mScale = new FloatProp(1f, endScale, DECELERATE_1_7);
+            FloatProp mAlpha = new FloatProp(1f, 0f, alphaInterpolator);
             FloatProp mShadowRadius = new FloatProp(startShadowRadius, 0, DECELERATE_1_7);
 
             @Override
@@ -1532,6 +1544,13 @@
         return closingAnimator;
     }
 
+    private boolean isFreeformAnimation(RemoteAnimationTarget[] appTargets) {
+        return DesktopModeStatus.canEnterDesktopMode(mLauncher.getApplicationContext())
+                && DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_EXIT_TRANSITIONS.isTrue()
+                && Arrays.stream(appTargets)
+                        .anyMatch(app -> app.taskInfo != null && app.taskInfo.isFreeform());
+    }
+
     private void addCujInstrumentation(Animator anim, int cuj) {
         anim.addListener(getCujAnimationSuccessListener(cuj));
     }
diff --git a/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java b/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java
index 92d9516..8e80aa5 100644
--- a/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java
+++ b/quickstep/src/com/android/launcher3/appprediction/PredictionRowView.java
@@ -57,7 +57,7 @@
     // Vertical padding of the icon that contributes to the expected cell height.
     private final int mVerticalPadding;
     // Extra padding that is used in the top app rows (prediction and search) that is not used in
-    // the regular A-Z list. This only applies to single line label.
+    // the regular A-Z list.
     private final int mTopRowExtraHeight;
 
     // Helper to drawing the focus indicator.
@@ -140,7 +140,7 @@
         // is not enabled. Otherwise, the extra height will increase by just the textHeight.
         int extraHeight = (Flags.enableTwolineToggle() &&
                 LauncherPrefs.ENABLE_TWOLINE_ALLAPPS_TOGGLE.get(getContext()))
-                ? textHeight : mTopRowExtraHeight;
+                ? (textHeight + mTopRowExtraHeight) : mTopRowExtraHeight;
         totalHeight += extraHeight;
         return getVisibility() == GONE ? 0 : totalHeight + getPaddingTop() + getPaddingBottom();
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
index 23a5a27..3b7ad3e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
@@ -139,18 +139,42 @@
             @NonNull Set<Integer> taskIdsToExclude,
             boolean wasOpenedFromTaskbar) {
         if (mQuickSwitchViewController != null) {
-            if (!mQuickSwitchViewController.isCloseAnimationRunning()
-                    && mQuickSwitchViewController.wasOpenedFromTaskbar() == wasOpenedFromTaskbar) {
-                return;
-            }
+            if (!mQuickSwitchViewController.isCloseAnimationRunning()) {
+                if (mQuickSwitchViewController.wasOpenedFromTaskbar() == wasOpenedFromTaskbar) {
+                    return;
+                }
 
-            // Allow the KQS to be reopened during the close animation to make it more responsive.
-            // Similarly, if KQS was opened in different mode (from taskbar vs. keyboard event),
-            // close it so it can be reopened in the correct mode.
-            // TODO(b/368119679) Consider updating list of shown tasks in place, or at least reopen
-            // the view in the same vertical location.
-            closeQuickSwitchView(false);
+                // Relayout the KQS view instead of recreating a new one if it is the current
+                // trigger surface is different than the previous one.
+                final int currentFocusIndexOverride =
+                        currentFocusedIndex == -1 && !mControllerCallbacks.isFirstTaskRunning()
+                                ? 0 : currentFocusedIndex;
+
+                // Skip the task reload if the list is not changed.
+                if (!mModel.isTaskListValid(mTaskListChangeId) || !taskIdsToExclude.equals(
+                        mExcludedTaskIds)) {
+                    mExcludedTaskIds = taskIdsToExclude;
+                    mTaskListChangeId = mModel.getTasks((tasks) -> {
+                        processLoadedTasks(tasks, taskIdsToExclude);
+                        mQuickSwitchViewController.updateQuickSwitchView(
+                                mTasks,
+                                mNumHiddenTasks,
+                                currentFocusIndexOverride,
+                                mHasDesktopTask,
+                                mWasDesktopTaskFilteredOut);
+                    });
+                }
+
+                mQuickSwitchViewController.updateLayoutForSurface(wasOpenedFromTaskbar,
+                        currentFocusIndexOverride);
+                return;
+            } else {
+                // Allow the KQS to be reopened during the close animation to make it more
+                // responsive.
+                closeQuickSwitchView(false);
+            }
         }
+
         mOverlayContext = mControllers.taskbarOverlayController.requestWindow();
         if (Flags.taskbarOverflow()) {
             mOverlayContext.getDragLayer().addTouchController(this);
@@ -186,13 +210,7 @@
 
         mExcludedTaskIds = taskIdsToExclude;
         mTaskListChangeId = mModel.getTasks((tasks) -> {
-            mHasDesktopTask = false;
-            mWasDesktopTaskFilteredOut = false;
-            if (onDesktop) {
-                processLoadedTasksOnDesktop(tasks, taskIdsToExclude);
-            } else {
-                processLoadedTasks(tasks, taskIdsToExclude);
-            }
+            processLoadedTasks(tasks, taskIdsToExclude);
             // Check if the first task is running after the recents model has updated so that we use
             // the correct index.
             mQuickSwitchViewController.openQuickSwitchView(
@@ -213,6 +231,17 @@
     }
 
     private void processLoadedTasks(List<GroupTask> tasks, Set<Integer> taskIdsToExclude) {
+        mHasDesktopTask = false;
+        mWasDesktopTaskFilteredOut = false;
+        if (mControllers.taskbarDesktopModeController.getAreDesktopTasksVisible()) {
+            processLoadedTasksOnDesktop(tasks, taskIdsToExclude);
+        } else {
+            processLoadedTasksOutsideDesktop(tasks, taskIdsToExclude);
+        }
+    }
+
+    private void processLoadedTasksOutsideDesktop(List<GroupTask> tasks,
+            Set<Integer> taskIdsToExclude) {
         // Only store MAX_TASK tasks, from most to least recent
         Collections.reverse(tasks);
         mTasks = tasks.stream()
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchView.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchView.java
index 05d34b5..1967dfd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchView.java
@@ -201,6 +201,8 @@
             int currentFocusIndexOverride,
             @NonNull KeyboardQuickSwitchViewController.ViewCallbacks viewCallbacks,
             boolean useDesktopTaskView) {
+        mContent.removeAllViews();
+
         mViewCallbacks = viewCallbacks;
         Resources resources = context.getResources();
         Resources.Theme theme = context.getTheme();
@@ -333,11 +335,17 @@
         return closeAnimation;
     }
 
-    private void animateOpen(int currentFocusIndexOverride) {
+    protected void animateOpen(int currentFocusIndexOverride) {
         if (mOpenAnimation != null) {
             // Restart animation since currentFocusIndexOverride can change the initial scroll.
             mOpenAnimation.cancel();
         }
+
+        // Reset the alpha for the case where the KQS view is opened before.
+        setAlpha(0);
+        mScrollView.setAlpha(0);
+        mNoRecentItemsPane.setAlpha(0);
+
         mOpenAnimation = new AnimatorSet();
 
         Animator outlineAnimation = mOutlineAnimationProgress.animateToValue(1f);
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
index 390112e..985cc26 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
@@ -128,6 +128,23 @@
                 /* useDesktopTaskView= */ !onDesktop && hasDesktopTask);
     }
 
+    protected void updateQuickSwitchView(
+            @NonNull List<GroupTask> tasks,
+            int numHiddenTasks,
+            int currentFocusIndexOverride,
+            boolean hasDesktopTask,
+            boolean wasDesktopTaskFilteredOut) {
+        mWasDesktopTaskFilteredOut = wasDesktopTaskFilteredOut;
+        mKeyboardQuickSwitchView.applyLoadPlan(
+                mOverlayContext,
+                tasks,
+                numHiddenTasks,
+                /* updateTasks= */ true,
+                currentFocusIndexOverride,
+                mViewCallbacks,
+                /* useDesktopTaskView= */ !mOnDesktop && hasDesktopTask);
+    }
+
     protected void positionView(boolean wasOpenedFromTaskbar, boolean isTransientTaskbar) {
         if (!wasOpenedFromTaskbar) {
             // Keep the default positioning.
@@ -155,6 +172,20 @@
         mKeyboardQuickSwitchView.setLayoutParams(lp);
     }
 
+    protected void updateLayoutForSurface(boolean updateLayoutFromTaskbar,
+            int currentFocusIndexOverride) {
+        BaseDragLayer.LayoutParams lp =
+                (BaseDragLayer.LayoutParams) mKeyboardQuickSwitchView.getLayoutParams();
+
+        if (updateLayoutFromTaskbar) {
+            lp.width = BaseDragLayer.LayoutParams.WRAP_CONTENT;
+        } else {
+            lp.width = BaseDragLayer.LayoutParams.MATCH_PARENT;
+        }
+
+        mKeyboardQuickSwitchView.animateOpen(currentFocusIndexOverride);
+    }
+
     boolean isCloseAnimationRunning() {
         return mCloseAnimation != null;
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index 4a94be7..f92d9dd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -33,6 +33,7 @@
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Flags;
+import com.android.launcher3.Hotseat;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimatedFloat;
@@ -83,6 +84,7 @@
     private final DeviceProfile.OnDeviceProfileChangeListener mOnDeviceProfileChangeListener =
             dp -> {
                 onStashedInAppChanged(dp);
+                adjustHotseatForBubbleBar();
                 if (mControllers != null && mControllers.taskbarViewController != null) {
                     mControllers.taskbarViewController.onRotationChanged(dp);
                 }
@@ -263,6 +265,14 @@
         }
     }
 
+    private void adjustHotseatForBubbleBar() {
+        Hotseat hotseat = mLauncher.getHotseat();
+        if (mControllers.bubbleControllers.isEmpty() || hotseat == null) return;
+        boolean hiddenForBubbles =
+                mControllers.bubbleControllers.get().bubbleBarViewController.isHiddenForNoBubbles();
+        hotseat.post(() -> adjustHotseatForBubbleBar(!hiddenForBubbles));
+    }
+
     /**
      * Create Taskbar animation when going from an app to Launcher as part of recents transition.
      * @param toState If known, the state we will end up in when reaching Launcher.
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
index 5a63ca6..db70724 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -221,10 +221,13 @@
         uiController = newUiController;
         uiController.init(this);
         uiController.updateStateForSysuiFlags(mSharedState.sysuiStateFlags);
-        // if bubble controllers are present take bubble bar location, else set it to null
+        // if bubble controllers are present configure the UI controller
         bubbleControllers.ifPresentOrElse(bubbleControllers -> {
             BubbleBarLocation location =
                     bubbleControllers.bubbleBarViewController.getBubbleBarLocation();
+            boolean hiddenForBubbles =
+                    bubbleControllers.bubbleBarViewController.isHiddenForNoBubbles();
+            uiController.adjustHotseatForBubbleBar(!hiddenForBubbles);
             uiController.onBubbleBarLocationUpdated(location);
         }, () -> uiController.onBubbleBarLocationUpdated(null));
         // Notify that the ui controller has changed
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 97d7179..31aa489 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -667,7 +667,7 @@
         TopTaskTracker.CachedTaskInfo cachedTaskInfo = mGestureState.getRunningTask();
         if (mIsSwipeForSplit) {
             int[] splitTaskIds = TopTaskTracker.INSTANCE.get(mContext).getRunningSplitTaskIds();
-            runningTasks = cachedTaskInfo.getPlaceholderTasks(splitTaskIds);
+            runningTasks = cachedTaskInfo.getSplitPlaceholderTasks(splitTaskIds);
         } else {
             runningTasks = cachedTaskInfo.getPlaceholderTasks();
         }
diff --git a/quickstep/src/com/android/quickstep/GestureState.java b/quickstep/src/com/android/quickstep/GestureState.java
index cff352c..5190ec8 100644
--- a/quickstep/src/com/android/quickstep/GestureState.java
+++ b/quickstep/src/com/android/quickstep/GestureState.java
@@ -28,6 +28,7 @@
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.SET_END_TARGET_HOME;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.SET_END_TARGET_NEW_TASK;
 
+import android.app.TaskInfo;
 import android.content.Intent;
 import android.os.SystemClock;
 import android.view.MotionEvent;
@@ -45,6 +46,7 @@
 import com.android.quickstep.util.ActiveGestureProtoLogProxy;
 import com.android.quickstep.views.RecentsViewContainer;
 import com.android.systemui.shared.recents.model.ThumbnailData;
+import com.android.wm.shell.shared.GroupedTaskInfo;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -330,13 +332,23 @@
         if (mRunningTask == null) {
             return new int[]{INVALID_TASK_ID, INVALID_TASK_ID};
         } else {
-            int cachedTasksSize = mRunningTask.mAllCachedTasks.size();
-            int count = Math.min(cachedTasksSize, getMultipleTasks ? 2 : 1);
-            int[] runningTaskIds = new int[count];
-            for (int i = 0; i < count; i++) {
-                runningTaskIds[i] = mRunningTask.mAllCachedTasks.get(i).taskId;
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                if (mRunningTask.getVisibleTasks().isEmpty()) {
+                    return new int[0];
+                }
+                GroupedTaskInfo topRunningTask = mRunningTask.getVisibleTasks().getFirst();
+                List<TaskInfo> groupedTasks = topRunningTask.getTaskInfoList();
+                return groupedTasks.stream().mapToInt(
+                        groupedTask -> groupedTask.taskId).toArray();
+            } else {
+                int cachedTasksSize = mRunningTask.mAllCachedTasks.size();
+                int count = Math.min(cachedTasksSize, getMultipleTasks ? 2 : 1);
+                int[] runningTaskIds = new int[count];
+                for (int i = 0; i < count; i++) {
+                    runningTaskIds[i] = mRunningTask.mAllCachedTasks.get(i).taskId;
+                }
+                return runningTaskIds;
             }
-            return runningTaskIds;
         }
     }
 
diff --git a/quickstep/src/com/android/quickstep/RecentTasksList.java b/quickstep/src/com/android/quickstep/RecentTasksList.java
index 714838a..85e2b6e 100644
--- a/quickstep/src/com/android/quickstep/RecentTasksList.java
+++ b/quickstep/src/com/android/quickstep/RecentTasksList.java
@@ -114,12 +114,9 @@
             }
 
             @Override
-            public void onTaskMovedToFront(GroupedTaskInfo[] visibleTasks) {
+            public void onTaskMovedToFront(GroupedTaskInfo taskToFront) {
                 mMainThreadExecutor.execute(() -> {
-                    // TODO(b/346588978): We currently are only sending a single task, but this will
-                    //                    be updated once we send the full set of visible tasks
-                    final TaskInfo info = visibleTasks[0].getTaskInfo1();
-                    topTaskTracker.handleTaskMovedToFront(info);
+                    topTaskTracker.handleTaskMovedToFront(taskToFront.getTaskInfo1());
                 });
             }
 
@@ -127,6 +124,13 @@
             public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
                 mMainThreadExecutor.execute(() -> topTaskTracker.onTaskChanged(taskInfo));
             }
+
+            @Override
+            public void onVisibleTasksChanged(GroupedTaskInfo[] visibleTasks) {
+                mMainThreadExecutor.execute(() -> {
+                    topTaskTracker.onVisibleTasksChanged(visibleTasks);
+                });
+            }
         });
         // We may receive onRunningTaskAppeared events later for tasks which have already been
         // included in the list returned by mSysUiProxy.getRunningTasks(), or may receive
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
index de8be50..e296449 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
@@ -354,7 +354,11 @@
      * @return whether the given running task info matches the gesture-blocked task.
      */
     public boolean isGestureBlockedTask(CachedTaskInfo taskInfo) {
-        return taskInfo != null && taskInfo.getTaskId() == mGestureBlockingTaskId;
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return taskInfo != null && taskInfo.topGroupedTaskContainsTask(mGestureBlockingTaskId);
+        } else {
+            return taskInfo != null && taskInfo.getTaskId() == mGestureBlockingTaskId;
+        }
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index c9dfe6d..80d6137 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -18,17 +18,18 @@
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_RECENTS;
-import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.content.Intent.ACTION_CHOOSER;
 import static android.content.Intent.FLAG_ACTIVITY_EXCLUDE_FROM_RECENTS;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT;
+import static com.android.wm.shell.shared.GroupedTaskInfo.TYPE_SPLIT;
 
-import android.annotation.UserIdInt;
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.TaskInfo;
 import android.content.Context;
+import android.util.ArrayMap;
+import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -46,8 +47,10 @@
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.shared.system.TaskStackChangeListeners;
+import com.android.wm.shell.shared.GroupedTaskInfo;
 import com.android.wm.shell.splitscreen.ISplitScreenListener;
 
+import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -61,37 +64,58 @@
  */
 public class TopTaskTracker extends ISplitScreenListener.Stub
         implements TaskStackChangeListener, SafeCloseable {
-
+    private static final String TAG = "TopTaskTracker";
     public static MainThreadInitializedObject<TopTaskTracker> INSTANCE =
             new MainThreadInitializedObject<>(TopTaskTracker::new);
 
     private static final int HISTORY_SIZE = 5;
 
+    private final Context mContext;
+
+    // Only used when Flags.enableShellTopTaskTracking() is disabled
     // Ordered list with first item being the most recent task.
     private final LinkedList<TaskInfo> mOrderedTaskList = new LinkedList<>();
-
-    private final Context mContext;
     private final SplitStageInfo mMainStagePosition = new SplitStageInfo();
     private final SplitStageInfo mSideStagePosition = new SplitStageInfo();
     private int mPinnedTaskId = INVALID_TASK_ID;
 
+    // Only used when Flags.enableShellTopTaskTracking() is enabled
+    // Mapping of display id to running tasks.  Running tasks are ordered from top most to
+    // bottom most.
+    private ArrayMap<Integer, ArrayList<GroupedTaskInfo>> mVisibleTasks = new ArrayMap<>();
+
     private TopTaskTracker(Context context) {
         mContext = context;
-        mMainStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_MAIN;
-        mSideStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_SIDE;
 
-        TaskStackChangeListeners.getInstance().registerTaskStackListener(this);
-        SystemUiProxy.INSTANCE.get(context).registerSplitScreenListener(this);
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            // Just prepopulate a list for the default display tasks so we don't need to add null
+            // checks everywhere
+            mVisibleTasks.put(DEFAULT_DISPLAY, new ArrayList<>());
+        } else {
+            mMainStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_MAIN;
+            mSideStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_SIDE;
+
+            TaskStackChangeListeners.getInstance().registerTaskStackListener(this);
+            SystemUiProxy.INSTANCE.get(context).registerSplitScreenListener(this);
+        }
     }
 
     @Override
     public void close() {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         TaskStackChangeListeners.getInstance().unregisterTaskStackListener(this);
         SystemUiProxy.INSTANCE.get(mContext).unregisterSplitScreenListener(this);
     }
 
     @Override
     public void onTaskRemoved(int taskId) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         mOrderedTaskList.removeIf(rto -> rto.taskId == taskId);
     }
 
@@ -100,7 +124,11 @@
         handleTaskMovedToFront(taskInfo);
     }
 
-    public void handleTaskMovedToFront(TaskInfo taskInfo) {
+    void handleTaskMovedToFront(TaskInfo taskInfo) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         mOrderedTaskList.removeIf(rto -> rto.taskId == taskInfo.taskId);
         mOrderedTaskList.addFirst(taskInfo);
 
@@ -131,8 +159,39 @@
         }
     }
 
+    /**
+     * Called when the set of visible tasks have changed.
+     */
+    public void onVisibleTasksChanged(GroupedTaskInfo[] visibleTasks) {
+        if (!com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
+        // TODO(346588978): Per-display info, just have everything in order by display
+
+        // Clear existing tasks for each display
+        mVisibleTasks.forEach((displayId, visibleTasksOnDisplay) -> visibleTasksOnDisplay.clear());
+
+        // Update the visible tasks on each display
+        for (int i = 0; i < visibleTasks.length; i++) {
+            final int displayId = visibleTasks[i].getTaskInfo1().getDisplayId();
+            final ArrayList<GroupedTaskInfo> displayTasks;
+            if (mVisibleTasks.containsKey(displayId)) {
+                displayTasks = mVisibleTasks.get(displayId);
+            } else {
+                displayTasks = new ArrayList<>();
+                mVisibleTasks.put(displayId, displayTasks);
+            }
+            displayTasks.add(visibleTasks[i]);
+        }
+    }
+
     @Override
     public void onStagePositionChanged(@StageType int stage, @StagePosition int position) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         if (stage == SplitConfigurationOptions.STAGE_TYPE_MAIN) {
             mMainStagePosition.stagePosition = position;
         } else {
@@ -141,6 +200,10 @@
     }
 
     public void onTaskChanged(RunningTaskInfo taskInfo) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         for (int i = 0; i < mOrderedTaskList.size(); i++) {
             if (mOrderedTaskList.get(i).taskId == taskInfo.taskId) {
                 mOrderedTaskList.set(i, taskInfo);
@@ -151,6 +214,10 @@
 
     @Override
     public void onTaskStageChanged(int taskId, @StageType int stage, boolean visible) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         // If a task is not visible anymore or has been moved to undefined, stop tracking it.
         if (!visible || stage == SplitConfigurationOptions.STAGE_TYPE_UNDEFINED) {
             if (mMainStagePosition.taskId == taskId) {
@@ -170,11 +237,19 @@
 
     @Override
     public void onActivityPinned(String packageName, int userId, int taskId, int stackId) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         mPinnedTaskId = taskId;
     }
 
     @Override
     public void onActivityUnpinned() {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         mPinnedTaskId = INVALID_TASK_ID;
     }
 
@@ -183,21 +258,59 @@
      * Will return empty array if device is not in staged split
      */
     public int[] getRunningSplitTaskIds() {
-        if (mMainStagePosition.taskId == INVALID_TASK_ID
-                || mSideStagePosition.taskId == INVALID_TASK_ID) {
-            return new int[]{};
-        }
-        int[] out = new int[2];
-        if (mMainStagePosition.stagePosition == STAGE_POSITION_TOP_OR_LEFT) {
-            out[0] = mMainStagePosition.taskId;
-            out[1] = mSideStagePosition.taskId;
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            // TODO(346588978): This assumes default display for now
+            final ArrayList<GroupedTaskInfo> visibleTasks = mVisibleTasks.get(DEFAULT_DISPLAY);
+            final GroupedTaskInfo splitTaskInfo = visibleTasks.stream()
+                    .filter(taskInfo -> taskInfo.getType() == TYPE_SPLIT)
+                    .findFirst().orElse(null);
+            if (splitTaskInfo != null && splitTaskInfo.getSplitBounds() != null) {
+                return new int[] {
+                        splitTaskInfo.getSplitBounds().leftTopTaskId,
+                        splitTaskInfo.getSplitBounds().rightBottomTaskId
+                };
+            }
+            return new int[0];
         } else {
-            out[1] = mMainStagePosition.taskId;
-            out[0] = mSideStagePosition.taskId;
+            if (mMainStagePosition.taskId == INVALID_TASK_ID
+                    || mSideStagePosition.taskId == INVALID_TASK_ID) {
+                return new int[]{};
+            }
+            int[] out = new int[2];
+            if (mMainStagePosition.stagePosition == STAGE_POSITION_TOP_OR_LEFT) {
+                out[0] = mMainStagePosition.taskId;
+                out[1] = mSideStagePosition.taskId;
+            } else {
+                out[1] = mMainStagePosition.taskId;
+                out[0] = mSideStagePosition.taskId;
+            }
+            return out;
         }
-        return out;
     }
 
+    /**
+     * Dumps the list of tasks in top task tracker.
+     */
+    public void dump(PrintWriter pw) {
+        if (!com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
+        // TODO(346588978): This assumes default display for now
+        final ArrayList<GroupedTaskInfo> displayTasks = mVisibleTasks.get(DEFAULT_DISPLAY);
+        pw.println("TopTaskTracker:");
+        pw.println("  tasks: [");
+        for (GroupedTaskInfo taskInfo : displayTasks) {
+            final TaskInfo info = taskInfo.getTaskInfo1();
+            final boolean isExcluded = (info.baseIntent.getFlags()
+                    & FLAG_ACTIVITY_EXCLUDE_FROM_RECENTS) != 0;
+            pw.println("    " + info.taskId + ": excluded=" + isExcluded
+                    + " visibleRequested=" + info.isVisibleRequested
+                    + " visible=" + info.isVisible
+                    + " " + info.baseIntent.getComponent());
+        }
+        pw.println("  ]");
+    }
 
     /**
      * Returns the CachedTaskInfo for the top most task
@@ -205,25 +318,35 @@
     @NonNull
     @UiThread
     public CachedTaskInfo getCachedTopTask(boolean filterOnlyVisibleRecents) {
-        if (filterOnlyVisibleRecents) {
-            // Since we only know about the top most task, any filtering may not be applied on the
-            // cache. The second to top task may change while the top task is still the same.
-            RunningTaskInfo[] tasks = TraceHelper.allowIpcs("getCachedTopTask.true", () ->
-                    ActivityManagerWrapper.getInstance().getRunningTasks(true));
-            return new CachedTaskInfo(Arrays.asList(tasks));
-        }
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            // TODO(346588978): Currently ignore filterOnlyVisibleRecents, but perhaps make this an
+            //  explicit filter For things to ignore (ie. PIP/Bubbles/Assistant/etc/so that this is
+            //  explicit)
+            // TODO(346588978): This assumes default display for now (as does all of Launcher)
+            final ArrayList<GroupedTaskInfo> displayTasks = mVisibleTasks.get(DEFAULT_DISPLAY);
+            return new CachedTaskInfo(new ArrayList<>(displayTasks));
+        } else {
+            if (filterOnlyVisibleRecents) {
+                // Since we only know about the top most task, any filtering may not be applied on
+                // the cache. The second to top task may change while the top task is still the
+                // same.
+                RunningTaskInfo[] tasks = TraceHelper.allowIpcs("getCachedTopTask.true", () ->
+                        ActivityManagerWrapper.getInstance().getRunningTasks(true));
+                return new CachedTaskInfo(Arrays.asList(tasks));
+            }
 
-        if (mOrderedTaskList.isEmpty()) {
-            RunningTaskInfo[] tasks = TraceHelper.allowIpcs("getCachedTopTask.false", () ->
-                    ActivityManagerWrapper.getInstance().getRunningTasks(
-                            false /* filterOnlyVisibleRecents */));
-            Collections.addAll(mOrderedTaskList, tasks);
-        }
+            if (mOrderedTaskList.isEmpty()) {
+                RunningTaskInfo[] tasks = TraceHelper.allowIpcs("getCachedTopTask.false", () ->
+                        ActivityManagerWrapper.getInstance().getRunningTasks(
+                                false /* filterOnlyVisibleRecents */));
+                Collections.addAll(mOrderedTaskList, tasks);
+            }
 
-        ArrayList<TaskInfo> tasks = new ArrayList<>(mOrderedTaskList);
-        // Strip the pinned task and recents task
-        tasks.removeIf(t -> t.taskId == mPinnedTaskId || isRecentsTask(t));
-        return new CachedTaskInfo(tasks);
+            ArrayList<TaskInfo> tasks = new ArrayList<>(mOrderedTaskList);
+            // Strip the pinned task and recents task
+            tasks.removeIf(t -> t.taskId == mPinnedTaskId || isRecentsTask(t));
+            return new CachedTaskInfo(tasks);
+        }
     }
 
     private static boolean isRecentsTask(TaskInfo task) {
@@ -237,24 +360,79 @@
      */
     public static class CachedTaskInfo {
 
+        // Only used when enableShellTopTaskTracking() is disabled
         @Nullable
         private final TaskInfo mTopTask;
+        @Nullable
         public final List<TaskInfo> mAllCachedTasks;
 
-        CachedTaskInfo(List<TaskInfo> allCachedTasks) {
+        // Only used when enableShellTopTaskTracking() is enabled
+        @Nullable
+        private final GroupedTaskInfo mTopGroupedTask;
+        @Nullable
+        private final ArrayList<GroupedTaskInfo> mVisibleTasks;
+
+
+        // Only used when enableShellTopTaskTracking() is enabled
+        CachedTaskInfo(@NonNull ArrayList<GroupedTaskInfo> visibleTasks) {
+            mAllCachedTasks = null;
+            mTopTask = null;
+            mVisibleTasks = visibleTasks;
+            mTopGroupedTask = !mVisibleTasks.isEmpty() ? mVisibleTasks.getFirst() : null;
+
+        }
+
+        // Only used when enableShellTopTaskTracking() is disabled
+        CachedTaskInfo(@NonNull List<TaskInfo> allCachedTasks) {
+            mVisibleTasks = null;
+            mTopGroupedTask = null;
             mAllCachedTasks = allCachedTasks;
             mTopTask = allCachedTasks.isEmpty() ? null : allCachedTasks.get(0);
         }
 
+        /**
+         * @return The list of visible tasks
+         */
+        public ArrayList<GroupedTaskInfo> getVisibleTasks() {
+            return mVisibleTasks;
+        }
+
+        /**
+         * @return The top task id
+         */
         public int getTaskId() {
-            return mTopTask == null ? INVALID_TASK_ID : mTopTask.taskId;
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // Callers should use topGroupedTaskContainsTask() instead
+                return INVALID_TASK_ID;
+            } else {
+                return mTopTask != null ? mTopTask.taskId : INVALID_TASK_ID;
+            }
+        }
+
+        /**
+         * @return Whether the top grouped task contains the given {@param taskId} if
+         *         Flags.enableShellTopTaskTracking() is true, otherwise it checks the top
+         *         task as reported from TaskStackListener.
+         */
+        public boolean topGroupedTaskContainsTask(int taskId) {
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                return mTopGroupedTask != null && mTopGroupedTask.containsTask(taskId);
+            } else {
+                return mTopTask != null && mTopTask.taskId == taskId;
+            }
         }
 
         /**
          * Returns true if the root of the task chooser activity
          */
         public boolean isRootChooseActivity() {
-            return mTopTask != null && ACTION_CHOOSER.equals(mTopTask.baseIntent.getAction());
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to not make an assumption on a specific task info
+                return mTopGroupedTask != null && ACTION_CHOOSER.equals(
+                        mTopGroupedTask.getTaskInfo1().baseIntent.getAction());
+            } else {
+                return mTopTask != null && ACTION_CHOOSER.equals(mTopTask.baseIntent.getAction());
+            }
         }
 
         /**
@@ -262,6 +440,10 @@
          * is another running task that is not excluded from recents, returns that underlying task.
          */
         public @Nullable CachedTaskInfo getVisibleNonExcludedTask() {
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // Callers should not need this when the full set of visible tasks are provided
+                return null;
+            }
             if (mTopTask == null
                     || (mTopTask.baseIntent.getFlags() & FLAG_ACTIVITY_EXCLUDE_FROM_RECENTS) == 0) {
                 // Not an excluded task.
@@ -278,24 +460,30 @@
         }
 
         /**
-         * Returns true if this represents the HOME task
+         * Returns true if this represents the HOME activity type task
          */
         public boolean isHomeTask() {
-            return mTopTask != null && mTopTask.configuration.windowConfiguration
-                    .getActivityType() == ACTIVITY_TYPE_HOME;
-        }
-
-        public boolean isRecentsTask() {
-            return TopTaskTracker.isRecentsTask(mTopTask);
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to not make an assumption on a specific task info
+                return mTopGroupedTask != null
+                        && mTopGroupedTask.getTaskInfo1().getActivityType() == ACTIVITY_TYPE_HOME;
+            } else {
+                return mTopTask != null && mTopTask.configuration.windowConfiguration
+                        .getActivityType() == ACTIVITY_TYPE_HOME;
+            }
         }
 
         /**
-         * Returns {@code true} if this task windowing mode is set to {@link
-         * android.app.WindowConfiguration#WINDOWING_MODE_FREEFORM}
+         * Returns true if this represents the RECENTS activity type task
          */
-        public boolean isFreeformTask() {
-            return mTopTask != null && mTopTask.configuration.windowConfiguration.getWindowingMode()
-                    == WINDOWING_MODE_FREEFORM;
+        public boolean isRecentsTask() {
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to not make an assumption on a specific task info
+                return mTopGroupedTask != null
+                        && TopTaskTracker.isRecentsTask(mTopGroupedTask.getTaskInfo1());
+            } else {
+                return TopTaskTracker.isRecentsTask(mTopTask);
+            }
         }
 
         /**
@@ -303,43 +491,78 @@
          * is loaded by the model
          */
         public Task[] getPlaceholderTasks() {
-            return mTopTask == null ? new Task[0]
-                    : new Task[]{Task.from(new TaskKey(mTopTask), mTopTask, false)};
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to return more than a single task once the callers
+                //  are refactored
+                if (mVisibleTasks.isEmpty()) {
+                    return new Task[0];
+                }
+                final TaskInfo info = mVisibleTasks.getFirst().getTaskInfo1();
+                return new Task[]{Task.from(new TaskKey(info), info, false)};
+            } else {
+                return mTopTask == null ? new Task[0]
+                        : new Task[]{Task.from(new TaskKey(mTopTask), mTopTask, false)};
+            }
         }
 
         /**
          * Returns {@link Task} array corresponding to the provided task ids which can be used as a
          * placeholder until the true object is loaded by the model
          */
-        public Task[] getPlaceholderTasks(int[] taskIds) {
-            if (mTopTask == null) {
-                return new Task[0];
-            }
-            Task[] result = new Task[taskIds.length];
-            for (int i = 0; i < taskIds.length; i++) {
-                final int index = i;
-                int taskId = taskIds[i];
-                mAllCachedTasks.forEach(rti -> {
-                    if (rti.taskId == taskId) {
-                        result[index] = Task.from(new TaskKey(rti), rti, false);
-                    }
-                });
-            }
-            return result;
-        }
+        public Task[] getSplitPlaceholderTasks(int[] taskIds) {
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                if (mVisibleTasks.isEmpty()
+                        || mVisibleTasks.getFirst().getType() != TYPE_SPLIT) {
+                    return new Task[0];
+                }
 
-        @UserIdInt
-        @Nullable
-        public Integer getUserId() {
-            return mTopTask == null ? null : mTopTask.userId;
+                GroupedTaskInfo splitTask = mVisibleTasks.getFirst();
+                Task[] result = new Task[taskIds.length];
+                for (int i = 0; i < taskIds.length; i++) {
+                    TaskInfo info = splitTask.getTaskById(taskIds[i]);
+                    if (info == null) {
+                        Log.w(TAG, "Requested task (" + taskIds[i] + ") not found");
+                        return new Task[0];
+                    }
+                    result[i] = Task.from(new TaskKey(info), info, false);
+                }
+                return result;
+            } else {
+                if (mTopTask == null) {
+                    return new Task[0];
+                }
+                Task[] result = new Task[taskIds.length];
+                for (int i = 0; i < taskIds.length; i++) {
+                    final int index = i;
+                    int taskId = taskIds[i];
+                    mAllCachedTasks.forEach(rti -> {
+                        if (rti.taskId == taskId) {
+                            result[index] = Task.from(new TaskKey(rti), rti, false);
+                        }
+                    });
+                }
+                return result;
+            }
         }
 
         @Nullable
         public String getPackageName() {
-            if (mTopTask == null || mTopTask.baseActivity == null) {
-                return null;
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to not make an assumption on a specific task info
+                if (mTopGroupedTask == null) {
+                    return null;
+                }
+                final TaskInfo info = mTopGroupedTask.getTaskInfo1();
+                if (info.baseActivity == null) {
+                    return null;
+                }
+                return info.baseActivity.getPackageName();
+            } else {
+                if (mTopTask == null || mTopTask.baseActivity == null) {
+                    return null;
+                }
+                return mTopTask.baseActivity.getPackageName();
             }
-            return mTopTask.baseActivity.getPackageName();
         }
     }
 }
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index ad5720f..0242fb6 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -1362,15 +1362,17 @@
                 && runningTask != null
                 && runningTask.isRootChooseActivity();
 
-        // In the case where we are in an excluded, translucent overlay, ignore it and treat the
-        // running activity as the task behind the overlay.
-        TopTaskTracker.CachedTaskInfo otherVisibleTask = runningTask == null
-                ? null
-                : runningTask.getVisibleNonExcludedTask();
-        if (otherVisibleTask != null) {
-            ActiveGestureProtoLogProxy.logUpdateGestureStateRunningTask(
-                    otherVisibleTask.getPackageName(), runningTask.getPackageName());
-            gestureState.updateRunningTask(otherVisibleTask);
+        if (!com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            // In the case where we are in an excluded, translucent overlay, ignore it and treat the
+            // running activity as the task behind the overlay.
+            TopTaskTracker.CachedTaskInfo otherVisibleTask = runningTask == null
+                    ? null
+                    : runningTask.getVisibleNonExcludedTask();
+            if (otherVisibleTask != null) {
+                ActiveGestureProtoLogProxy.logUpdateGestureStateRunningTask(
+                        otherVisibleTask.getPackageName(), runningTask.getPackageName());
+                gestureState.updateRunningTask(otherVisibleTask);
+            }
         }
 
         boolean previousGestureAnimatedToLauncher =
@@ -1672,6 +1674,7 @@
         ContextualSearchStateManager.INSTANCE.get(this).dump("\t", pw);
         SystemUiProxy.INSTANCE.get(this).dump(pw);
         DeviceConfigWrapper.get().dump("   ", pw);
+        TopTaskTracker.INSTANCE.get(this).dump(pw);
     }
 
     private AbsSwipeUpHandler createLauncherSwipeHandler(
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java
index 4f9d837..c1d3f6e 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowSwipeHandler.java
@@ -149,7 +149,8 @@
         if (mActiveAnimationFactory != null) {
             return;
         }
-        setHomeScaleAndAlpha(builder, app, mCurrentShift.value, 0);
+        setHomeScaleAndAlpha(builder, app, mCurrentShift.value,
+                Utilities.boundToRange(1 - mCurrentShift.value, 0, 1));
     }
 
     private void setHomeScaleAndAlpha(SurfaceProperties builder,
diff --git a/quickstep/src/com/android/quickstep/util/AppPairsController.java b/quickstep/src/com/android/quickstep/util/AppPairsController.java
index 7388d59..1312aa4 100644
--- a/quickstep/src/com/android/quickstep/util/AppPairsController.java
+++ b/quickstep/src/com/android/quickstep/util/AppPairsController.java
@@ -409,8 +409,8 @@
             );
         } else {
             // Tapped an app pair while in a single app
-            int runningTaskId = topTaskTracker
-                    .getCachedTopTask(false /* filterOnlyVisibleRecents */).getTaskId();
+            final TopTaskTracker.CachedTaskInfo runningTask = topTaskTracker
+                    .getCachedTopTask(false /* filterOnlyVisibleRecents */);
 
             mSplitSelectStateController.findLastActiveTasksAndRunCallback(
                     componentKeys,
@@ -418,10 +418,21 @@
                     foundTasks -> {
                         Task foundTask1 = foundTasks[0];
                         Task foundTask2 = foundTasks[1];
-                        boolean task1IsOnScreen =
-                                foundTask1 != null && foundTask1.getKey().getId() == runningTaskId;
-                        boolean task2IsOnScreen =
-                                foundTask2 != null && foundTask2.getKey().getId() == runningTaskId;
+                        boolean task1IsOnScreen;
+                        boolean task2IsOnScreen;
+                        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                            task1IsOnScreen = foundTask1 != null
+                                    && runningTask.topGroupedTaskContainsTask(
+                                    foundTask1.getKey().getId());
+                            task2IsOnScreen = foundTask2 != null
+                                    && runningTask.topGroupedTaskContainsTask(
+                                    foundTask2.getKey().getId());
+                        } else {
+                            task1IsOnScreen = foundTask1 != null && foundTask1.getKey().getId()
+                                    == runningTask.getTaskId();
+                            task2IsOnScreen = foundTask2 != null && foundTask2.getKey().getId()
+                                    == runningTask.getTaskId();
+                        }
 
                         if (!task1IsOnScreen && !task2IsOnScreen) {
                             // Neither App A nor App B are on-screen, launch the app pair normally.
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 3a4e328..8982850 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -847,6 +847,8 @@
     private final RecentsViewModelHelper mHelper;
     private final RecentsViewUtils mUtils = new RecentsViewUtils();
 
+    private final Matrix mTmpMatrix = new Matrix();
+
     public RecentsView(Context context, @Nullable AttributeSet attrs, int defStyleAttr,
             BaseContainerInterface sizeStrategy) {
         super(context, attrs, defStyleAttr);
@@ -1966,7 +1968,7 @@
                     // We try to avoid this because it can cause a scroll jump, but it is needed
                     // for cases where the running task isn't included in this load plan (e.g. if
                     // the current running task is excludedFromRecents.)
-                    showCurrentTask(mActiveGestureRunningTasks);
+                    showCurrentTask(mActiveGestureRunningTasks, "applyLoadPlan");
                 } else {
                     setRunningTaskViewId(INVALID_TASK_ID);
                 }
@@ -2749,7 +2751,7 @@
             updateSizeAndPadding();
         }
 
-        showCurrentTask(mActiveGestureRunningTasks);
+        showCurrentTask(mActiveGestureRunningTasks, "onGestureAnimationStart");
         setEnableFreeScroll(false);
         setEnableDrawingLiveTile(false);
         setRunningTaskHidden(true);
@@ -2930,8 +2932,9 @@
      * All subsequent calls to reload will keep the task as the first item until {@link #reset()}
      * is called.  Also scrolls the view to this task.
      */
-    private void showCurrentTask(Task[] runningTasks) {
-        Log.d(TAG, "showCurrentTask - runningTasks: " + Arrays.toString(runningTasks));
+    private void showCurrentTask(Task[] runningTasks, String caller) {
+        Log.d(TAG, "showCurrentTask(" + caller + ") - runningTasks: "
+                + Arrays.toString(runningTasks));
         if (runningTasks.length == 0) {
             return;
         }
@@ -5807,6 +5810,14 @@
         // mSyncTransactionApplier doesn't get transferred over
         runActionOnRemoteHandles(remoteTargetHandle -> {
             final TransformParams params = remoteTargetHandle.getTransformParams();
+            if (Flags.enableFallbackOverviewInWindow() || Flags.enableLauncherOverviewInWindow()) {
+                params.setHomeBuilderProxy((builder, app, transformParams) -> {
+                    mTmpMatrix.setScale(
+                            1f, 1f, app.localBounds.exactCenterX(), app.localBounds.exactCenterY());
+                    builder.setMatrix(mTmpMatrix).setAlpha(1f).setShow();
+                });
+            }
+
             if (mSyncTransactionApplier != null) {
                 params.setSyncTransactionApplier(mSyncTransactionApplier);
                 params.getTargetSet().addReleaseCheck(mSyncTransactionApplier);
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/AppPairsControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/AppPairsControllerTest.kt
index 541a48d..ee70e0a 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/AppPairsControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/AppPairsControllerTest.kt
@@ -202,7 +202,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -226,7 +226,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -250,7 +250,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -274,7 +274,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -298,7 +298,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -322,7 +322,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -341,12 +341,16 @@
         whenever(mockTaskKey1.getId()).thenReturn(1)
         whenever(mockTaskKey2.getId()).thenReturn(2)
         // ... with app 1 already on screen
-        whenever(mockCachedTaskInfo.taskId).thenReturn(1)
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(1))).thenReturn(true)
+        } else {
+            whenever(mockCachedTaskInfo.taskId).thenReturn(1)
+        }
 
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -365,12 +369,16 @@
         whenever(mockTaskKey1.getId()).thenReturn(1)
         whenever(mockTaskKey2.getId()).thenReturn(2)
         // ... with app 2 already on screen
-        whenever(mockCachedTaskInfo.taskId).thenReturn(2)
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(2))).thenReturn(true)
+        } else {
+            whenever(mockCachedTaskInfo.taskId).thenReturn(2)
+        }
 
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -389,12 +397,16 @@
         whenever(mockTaskKey1.getId()).thenReturn(1)
         whenever(mockTaskKey2.getId()).thenReturn(2)
         // ... with app 3 already on screen
-        whenever(mockCachedTaskInfo.taskId).thenReturn(3)
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(3))).thenReturn(true)
+        } else {
+            whenever(mockCachedTaskInfo.taskId).thenReturn(3)
+        }
 
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsKeyboardQuickSwitch.java b/quickstep/tests/src/com/android/quickstep/TaplTestsKeyboardQuickSwitch.java
index 43ebb17..3c4f1d9 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsKeyboardQuickSwitch.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsKeyboardQuickSwitch.java
@@ -49,6 +49,7 @@
         DISMISS(0),
         LAUNCH_LAST_APP(0),
         LAUNCH_SELECTED_APP(1),
+        DISMISS_WHEN_GOING_HOME(1),
         LAUNCH_OVERVIEW(KeyboardQuickSwitchController.MAX_TASKS - 1);
 
         private final int mNumAdditionalRunningTasks;
@@ -156,6 +157,11 @@
         mLauncher.goHome().showQuickSwitchView().launchFocusedAppTask(CALCULATOR_APP_PACKAGE);
     }
 
+    @Test
+    public void testDismissedWhenGoingHome() {
+        runTest(TestSurface.LAUNCHED_APP, TestCase.DISMISS_WHEN_GOING_HOME);
+    }
+
     private void runTest(@NonNull TestSurface testSurface, @NonNull TestCase testCase) {
         for (int i = 0; i < testCase.mNumAdditionalRunningTasks; i++) {
             startTestActivity(3 + i);
@@ -197,6 +203,9 @@
                 }
                 kqs.launchFocusedAppTask(CALCULATOR_APP_PACKAGE);
                 break;
+            case DISMISS_WHEN_GOING_HOME:
+                kqs.dismissByGoingHome();
+                break;
             case LAUNCH_OVERVIEW:
                 kqs.moveFocusBackward();
                 if (!testSurface.mInitialFocusAtZero) {
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index 77d789f..4ccdd53 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -312,9 +312,10 @@
         <attr name="rowCountSpecsId" format="reference" />
         <!-- defaults to allAppsCellSpecsId, if not specified -->
         <attr name="allAppsCellSpecsTwoPanelId" format="reference" />
-
         <!-- defaults to false, if not specified -->
         <attr name="isFixedLandscape" format="boolean" />
+        <!-- defaults to false, if not specified -->
+        <attr name="isOldGrid" format="boolean" />
 
         <!-- By default all categories are enabled -->
         <attr name="deviceCategory" format="integer">
diff --git a/src/com/android/launcher3/DeleteDropTarget.java b/src/com/android/launcher3/DeleteDropTarget.java
index 58789fd..425f277 100644
--- a/src/com/android/launcher3/DeleteDropTarget.java
+++ b/src/com/android/launcher3/DeleteDropTarget.java
@@ -130,7 +130,6 @@
     public void completeDrop(DragObject d) {
         ItemInfo item = d.dragInfo;
         if (canRemove(item)) {
-            onAccessibilityDrop(null, item);
             mDropTargetHandler.onDeleteComplete(item);
         }
     }
diff --git a/src/com/android/launcher3/Hotseat.java b/src/com/android/launcher3/Hotseat.java
index 6468f74..b2ccba4 100644
--- a/src/com/android/launcher3/Hotseat.java
+++ b/src/com/android/launcher3/Hotseat.java
@@ -187,22 +187,20 @@
     public void adjustForBubbleBar(boolean isBubbleBarVisible) {
         DeviceProfile dp = mActivity.getDeviceProfile();
         float adjustedBorderSpace = dp.getHotseatAdjustedBorderSpaceForBubbleBar(getContext());
-        boolean adjustmentRequired = Float.compare(adjustedBorderSpace, 0f) != 0;
-
+        boolean shouldAdjustHotseat = isBubbleBarVisible
+                && Float.compare(adjustedBorderSpace, 0f) != 0;
         ShortcutAndWidgetContainer icons = getShortcutsAndWidgets();
         // update the translation provider for future layout passes of hotseat icons.
-        if (adjustmentRequired && isBubbleBarVisible) {
+        if (shouldAdjustHotseat) {
             icons.setTranslationProvider(
                     cellX -> dp.getHotseatAdjustedTranslation(getContext(), cellX));
         } else {
             icons.setTranslationProvider(null);
         }
-        if (!adjustmentRequired) return;
-
         AnimatorSet animatorSet = new AnimatorSet();
         for (int i = 0; i < icons.getChildCount(); i++) {
             View child = icons.getChildAt(i);
-            float tx = isBubbleBarVisible ? dp.getHotseatAdjustedTranslation(getContext(), i) : 0;
+            float tx = shouldAdjustHotseat ? dp.getHotseatAdjustedTranslation(getContext(), i) : 0;
             if (child instanceof Reorderable) {
                 MultiTranslateDelegate mtd = ((Reorderable) child).getTranslateDelegate();
                 animatorSet.play(
@@ -213,8 +211,8 @@
         }
         if (mQsb instanceof HorizontalInsettableView horizontalInsettableQsb) {
             final float currentInsetFraction = horizontalInsettableQsb.getHorizontalInsets();
-            final float targetInsetFraction =
-                    isBubbleBarVisible ? (float) dp.iconSizePx / dp.hotseatQsbWidth : 0;
+            final float targetInsetFraction = shouldAdjustHotseat
+                    ? (float) dp.iconSizePx / dp.hotseatQsbWidth : 0;
             ValueAnimator qsbAnimator =
                     ValueAnimator.ofFloat(currentInsetFraction, targetInsetFraction);
             qsbAnimator.addUpdateListener(animation -> {
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 04e4b57..e18862a 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -1031,6 +1031,7 @@
         private final int mAllAppsCellSpecsTwoPanelId;
         private final int mRowCountSpecsId;
         private final boolean mIsFixedLandscape;
+        private final boolean mIsOldGrid;
 
         public GridOption(Context context, AttributeSet attrs, Info displayInfo) {
             TypedArray a = context.obtainStyledAttributes(
@@ -1175,6 +1176,7 @@
             }
 
             mIsFixedLandscape = a.getBoolean(R.styleable.GridDisplayOption_isFixedLandscape, false);
+            mIsOldGrid = a.getBoolean(R.styleable.GridDisplayOption_isOldGrid, false);
 
             int inlineForRotation = a.getInt(R.styleable.GridDisplayOption_inlineQsb,
                     DONT_INLINE_QSB);
@@ -1206,20 +1208,30 @@
             }
         }
 
-        public boolean isNewGridOption() {
-            return mRowCountSpecsId != INVALID_RESOURCE_HANDLE;
-        }
-
+        /**
+         * Returns true if the grid option should be used given the flags that are toggled on/off.
+         */
         public boolean filterByFlag(int deviceType, boolean isFixedLandscape) {
             if (deviceType == TYPE_TABLET) {
                 return Flags.oneGridRotationHandling() == mIsDualGrid;
             }
 
-            if (isFixedLandscape) {
-                return Flags.oneGridSpecs() && mIsFixedLandscape;
+            // Here we return true if fixed landscape mode should be on.
+            if (mIsFixedLandscape || isFixedLandscape) {
+                return mIsFixedLandscape && isFixedLandscape && Flags.oneGridSpecs();
             }
 
-            return ((Flags.oneGridSpecs() == isNewGridOption()) && !mIsFixedLandscape);
+            // Here we return true if we want to show the new grids.
+            if (mRowCountSpecsId != INVALID_RESOURCE_HANDLE) {
+                return Flags.oneGridSpecs();
+            }
+
+            // Here we return true if we want to show the old grids.
+            if (mIsOldGrid) {
+                return !Flags.oneGridSpecs();
+            }
+
+            return true;
         }
     }
 
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 0e9c861..95dbf5f 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3;
 
+import static com.android.launcher3.AbstractFloatingView.TYPE_WIDGET_RESIZE_FRAME;
 import static com.android.launcher3.BubbleTextView.DISPLAY_FOLDER;
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
 import static com.android.launcher3.LauncherAnimUtils.SPRING_LOADED_EXIT_DELAY;
@@ -1222,6 +1223,10 @@
     }
 
     protected void onPageBeginTransition() {
+        // Widget resize frame doesn't receive events to close when talkback is enabled. For that
+        // case, close it here.
+        AbstractFloatingView.closeOpenViews(mLauncher, false, TYPE_WIDGET_RESIZE_FRAME);
+
         super.onPageBeginTransition();
         updateChildrenLayersEnabled();
     }
diff --git a/src/com/android/launcher3/dagger/LauncherComponentProvider.kt b/src/com/android/launcher3/dagger/LauncherComponentProvider.kt
new file mode 100644
index 0000000..5015e54
--- /dev/null
+++ b/src/com/android/launcher3/dagger/LauncherComponentProvider.kt
@@ -0,0 +1,57 @@
+/*
+ * 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.dagger
+
+import android.content.Context
+import android.view.LayoutInflater
+import com.android.launcher3.LauncherApplication
+
+/**
+ * Utility class to extract LauncherAppComponent from a context.
+ *
+ * If the context doesn't provide LauncherAppComponent by default, it creates a new one and
+ * associate it with that context
+ */
+object LauncherComponentProvider {
+
+    @JvmStatic
+    fun get(c: Context): LauncherAppComponent {
+        val app = c.applicationContext
+        if (app is LauncherApplication) return app.appComponent
+
+        val inflater = LayoutInflater.from(app)
+        val existingFilter = inflater.filter
+        if (existingFilter is Holder) return existingFilter.component
+
+        // Create a new component
+        return Holder(
+                DaggerLauncherAppComponent.builder().appContext(app).build()
+                    as LauncherAppComponent,
+                existingFilter,
+            )
+            .apply { inflater.filter = this }
+            .component
+    }
+
+    private data class Holder(
+        val component: LauncherAppComponent,
+        private val filter: LayoutInflater.Filter?,
+    ) : LayoutInflater.Filter {
+
+        override fun onLoadClass(clazz: Class<*>?) = filter?.onLoadClass(clazz) ?: true
+    }
+}
diff --git a/src/com/android/launcher3/dragndrop/LauncherDragController.java b/src/com/android/launcher3/dragndrop/LauncherDragController.java
index 29fc613..4aa3673 100644
--- a/src/com/android/launcher3/dragndrop/LauncherDragController.java
+++ b/src/com/android/launcher3/dragndrop/LauncherDragController.java
@@ -119,6 +119,9 @@
                         initialDragViewScale,
                         dragViewScaleOnDrop,
                         scalePx);
+        // During a drag, we don't want to expose the descendendants of drag view to a11y users,
+        // since those decendents are not a valid position in the workspace.
+        dragView.setImportantForAccessibility(View.IMPORTANT_FOR_ACCESSIBILITY_NO_HIDE_DESCENDANTS);
         dragView.setItemInfo(dragInfo);
         mDragObject.dragComplete = false;
 
diff --git a/src/com/android/launcher3/settings/SettingsActivity.java b/src/com/android/launcher3/settings/SettingsActivity.java
index 5851f62..5068b48 100644
--- a/src/com/android/launcher3/settings/SettingsActivity.java
+++ b/src/com/android/launcher3/settings/SettingsActivity.java
@@ -22,6 +22,7 @@
 
 import static com.android.launcher3.BuildConfig.IS_DEBUG_DEVICE;
 import static com.android.launcher3.BuildConfig.IS_STUDIO_BUILD;
+import static com.android.launcher3.InvariantDeviceProfile.TYPE_MULTI_DISPLAY;
 import static com.android.launcher3.states.RotationHelper.ALLOW_ROTATION_PREFERENCE_KEY;
 
 import android.app.Activity;
@@ -52,6 +53,7 @@
 
 import com.android.launcher3.BuildConfig;
 import com.android.launcher3.Flags;
+import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherFiles;
 import com.android.launcher3.R;
 import com.android.launcher3.states.RotationHelper;
@@ -310,7 +312,10 @@
                     }
                     return mDeveloperOptionsEnabled;
                 case FIXED_LANDSCAPE_MODE:
-                    if (!Flags.oneGridSpecs()) {
+                    if (!Flags.oneGridSpecs()
+                            // adding this condition until fixing b/378972567
+                            || InvariantDeviceProfile.INSTANCE.get(getContext()).deviceType
+                            == TYPE_MULTI_DISPLAY) {
                         return false;
                     }
                     // When the setting changes rotate the screen accordingly to showcase the result
diff --git a/src/com/android/launcher3/util/DaggerSingletonObject.java b/src/com/android/launcher3/util/DaggerSingletonObject.java
index febe6af..a245761 100644
--- a/src/com/android/launcher3/util/DaggerSingletonObject.java
+++ b/src/com/android/launcher3/util/DaggerSingletonObject.java
@@ -18,8 +18,8 @@
 
 import android.content.Context;
 
-import com.android.launcher3.LauncherApplication;
 import com.android.launcher3.dagger.LauncherAppComponent;
+import com.android.launcher3.dagger.LauncherComponentProvider;
 
 import java.util.function.Function;
 
@@ -37,8 +37,6 @@
     }
 
     public T get(Context context) {
-        LauncherAppComponent component =
-                ((LauncherApplication) context.getApplicationContext()).getAppComponent();
-        return mFunction.apply(component);
+        return mFunction.apply(LauncherComponentProvider.get(context));
     }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 8bebfb2..150806a 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -295,8 +295,11 @@
     }
 
     private void attachScrollbarToRecyclerView(WidgetsRecyclerView recyclerView) {
-        recyclerView.bindFastScrollbar(mFastScroller, WIDGET_SCROLLER);
         if (mCurrentWidgetsRecyclerView != recyclerView) {
+            // Bind scrollbar if changing the recycler view. If widgets list updates, since
+            // scrollbar is already attached to the recycler view, it will automatically adjust as
+            // needed with recycler view's onScrollListener.
+            recyclerView.bindFastScrollbar(mFastScroller, WIDGET_SCROLLER);
             // Only reset the scroll position & expanded apps if the currently shown recycler view
             // has been updated.
             reset();
diff --git a/tests/Android.bp b/tests/Android.bp
index 35a2275..b1d4ef6 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -173,6 +173,7 @@
         "multivalentTests/src/**/*.java",
         "multivalentTests/src/**/*.kt",
         "src/com/android/launcher3/ui/AbstractLauncherUiTest.java",
+        "src/com/android/launcher3/ui/BaseLauncherTaplTest.java",
         "tapl/com/android/launcher3/tapl/*.java",
         "tapl/com/android/launcher3/tapl/*.kt",
     ],
diff --git a/tests/multivalentTests/src/com/android/launcher3/dagger/LauncherComponentProviderTest.kt b/tests/multivalentTests/src/com/android/launcher3/dagger/LauncherComponentProviderTest.kt
new file mode 100644
index 0000000..9255877
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/dagger/LauncherComponentProviderTest.kt
@@ -0,0 +1,74 @@
+/*
+ * 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.dagger
+
+import android.content.Context
+import android.content.ContextWrapper
+import android.view.ContextThemeWrapper
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.R
+import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext
+import com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE
+import org.junit.Assert.assertNotNull
+import org.junit.Assert.assertNotSame
+import org.junit.Assert.assertSame
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class LauncherComponentProviderTest {
+
+    val app: Context = ApplicationProvider.getApplicationContext()
+
+    @Test
+    fun `returns same component as Launcher application`() {
+        val c = SandboxModelContext()
+        assertSame(c.appComponent, LauncherComponentProvider.get(c))
+        assertNotSame(LauncherComponentProvider.get(c), LauncherComponentProvider.get(app))
+    }
+
+    @Test
+    fun `returns same component for isolated context`() {
+        val c = IsolatedContext()
+
+        // Same component is returned for multiple calls, irrespective of the wrappers
+        assertNotNull(LauncherComponentProvider.get(c))
+        assertSame(
+            LauncherComponentProvider.get(c),
+            LauncherComponentProvider.get(ContextThemeWrapper(c, R.style.LauncherTheme)),
+        )
+
+        // Different than main application
+        assertNotSame(LauncherComponentProvider.get(c), LauncherComponentProvider.get(app))
+    }
+
+    @Test
+    fun `different components for different isolated context`() {
+        val c1 = IsolatedContext()
+        val c2 = IsolatedContext()
+
+        assertNotNull(LauncherComponentProvider.get(c1))
+        assertNotNull(LauncherComponentProvider.get(c2))
+        assertNotSame(LauncherComponentProvider.get(c1), LauncherComponentProvider.get(c2))
+    }
+
+    inner class IsolatedContext : ContextWrapper(app.createPackageContext(TEST_PACKAGE, 0)) {
+
+        override fun getApplicationContext(): Context = this
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 2b1fddc..ed5762d 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -15,72 +15,42 @@
  */
 package com.android.launcher3.ui;
 
-import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
-
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 
-import static com.android.launcher3.testing.shared.TestProtocol.ICON_MISSING;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.TestUtil.resolveSystemAppInfo;
 
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
-import android.app.ActivityManager;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
-import android.content.pm.PackageInfo;
-import android.content.pm.PackageManager;
-import android.graphics.Point;
-import android.os.Debug;
 import android.os.Process;
-import android.os.RemoteException;
-import android.os.UserHandle;
-import android.os.UserManager;
-import android.platform.test.flag.junit.SetFlagsRule;
-import android.platform.test.rule.LimitDevicesRule;
 import android.system.OsConstants;
 import android.util.Log;
 
-import androidx.annotation.NonNull;
-import androidx.test.InstrumentationRegistry;
 import androidx.test.uiautomator.By;
 import androidx.test.uiautomator.BySelector;
-import androidx.test.uiautomator.UiDevice;
 import androidx.test.uiautomator.Until;
 
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.Utilities;
-import com.android.launcher3.celllayout.FavoriteItemsTransaction;
-import com.android.launcher3.tapl.HomeAllApps;
-import com.android.launcher3.tapl.HomeAppIcon;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.tapl.TestHelpers;
 import com.android.launcher3.testcomponent.TestCommandReceiver;
 import com.android.launcher3.util.LooperExecutor;
 import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.Wait;
-import com.android.launcher3.util.rule.ExtendedLongPressTimeoutRule;
 import com.android.launcher3.util.rule.FailureWatcher;
-import com.android.launcher3.util.rule.SamplerRule;
-import com.android.launcher3.util.rule.ScreenRecordRule;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.util.rule.TestIsolationRule;
-import com.android.launcher3.util.rule.TestStabilityRule;
 import com.android.launcher3.util.rule.ViewCaptureRule;
 
-import org.junit.After;
-import org.junit.Assert;
-import org.junit.Before;
-import org.junit.Rule;
 import org.junit.rules.RuleChain;
 import org.junit.rules.TestRule;
 
-import java.io.IOException;
 import java.util.Objects;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
@@ -93,169 +63,51 @@
 /**
  * Base class for all instrumentation tests providing various utility methods.
  */
-public abstract class AbstractLauncherUiTest<LAUNCHER_TYPE extends Launcher> {
-
-    public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 10;
+public abstract class AbstractLauncherUiTest<LAUNCHER_TYPE extends Launcher>
+        extends BaseLauncherTaplTest {
 
     private static final String TAG = "AbstractLauncherUiTest";
 
-    private static final long BYTES_PER_MEGABYTE = 1 << 20;
-
-    private static boolean sDumpWasGenerated = false;
-    private static boolean sActivityLeakReported = false;
-    private static boolean sSeenKeyguard = false;
-    private static boolean sFirstTimeWaitingForWizard = true;
-
-    private static final String SYSTEMUI_PACKAGE = "com.android.systemui";
-
     protected LooperExecutor mMainThreadExecutor = MAIN_EXECUTOR;
-    protected final UiDevice mDevice = getUiDevice();
-    protected final LauncherInstrumentation mLauncher = createLauncherInstrumentation();
-
-    @NonNull
-    public static LauncherInstrumentation createLauncherInstrumentation() {
-        waitForSetupWizardDismissal(); // precondition for creating LauncherInstrumentation
-        return new LauncherInstrumentation(true);
-    }
-
-    protected Context mTargetContext;
-    protected String mTargetPackage;
-    private int mLauncherPid;
-
-    private final ActivityManager.MemoryInfo mMemoryInfo = new ActivityManager.MemoryInfo();
-    private final ActivityManager mActivityManager;
-    private long mMemoryBefore;
-
-    /** Detects activity leaks and throws an exception if a leak is found. */
-    public static void checkDetectedLeaks(LauncherInstrumentation launcher) {
-        checkDetectedLeaks(launcher, false);
-    }
-
-    /** Detects activity leaks and throws an exception if a leak is found. */
-    public static void checkDetectedLeaks(LauncherInstrumentation launcher,
-            boolean requireOneActiveActivityUnused) {
-        if (TestStabilityRule.isPresubmit()) return; // b/313501215
-
-        final boolean requireOneActiveActivity =
-                false; // workaround for leaks when there is an unexpected Recents activity
-
-        if (sActivityLeakReported) return;
-
-        // Check whether activity leak detector has found leaked activities.
-        Wait.atMost(() -> getActivityLeakErrorMessage(launcher, requireOneActiveActivity),
-                () -> {
-                    launcher.forceGc();
-                    return MAIN_EXECUTOR.submit(
-                            () -> launcher.noLeakedActivities(requireOneActiveActivity)).get();
-                }, launcher);
-    }
-
-    public static String getAppPackageName() {
-        return getInstrumentation().getContext().getPackageName();
-    }
-
-    private static String getActivityLeakErrorMessage(LauncherInstrumentation launcher,
-            boolean requireOneActiveActivity) {
-        sActivityLeakReported = true;
-        return "Activity leak detector has found leaked activities, requirining 1 activity: "
-                + requireOneActiveActivity + "; "
-                + dumpHprofData(launcher, false, requireOneActiveActivity) + ".";
-    }
-
-    private static String dumpHprofData(LauncherInstrumentation launcher, boolean intentionalLeak,
-            boolean requireOneActiveActivity) {
-        if (intentionalLeak) return "intentional leak; not generating dump";
-
-        String result;
-        if (sDumpWasGenerated) {
-            result = "dump has already been generated by another test";
-        } else {
-            try {
-                final String fileName =
-                        getInstrumentation().getTargetContext().getFilesDir().getPath()
-                                + "/ActivityLeakHeapDump.hprof";
-                if (TestHelpers.isInLauncherProcess()) {
-                    Debug.dumpHprofData(fileName);
-                } else {
-                    final UiDevice device = getUiDevice();
-                    device.executeShellCommand(
-                            "am dumpheap " + device.getLauncherPackageName() + " " + fileName);
-                }
-                Log.d(TAG, "Saved leak dump, the leak is still present: "
-                        + !launcher.noLeakedActivities(requireOneActiveActivity));
-                sDumpWasGenerated = true;
-                result = "saved memory dump as an artifact";
-            } catch (Throwable e) {
-                Log.e(TAG, "dumpHprofData failed", e);
-                result = "failed to save memory dump";
-            }
-        }
-        return result + ". Full list of activities: " + launcher.getRootedActivitiesList();
-    }
 
     protected AbstractLauncherUiTest() {
-        mActivityManager = InstrumentationRegistry.getContext()
-                .getSystemService(ActivityManager.class);
-        mLauncher.enableCheckEventsForSuccessfulGestures();
-        mLauncher.setAnomalyChecker(AbstractLauncherUiTest::verifyKeyguardInvisible);
-        try {
-            mDevice.setOrientationNatural();
-        } catch (RemoteException e) {
-            throw new RuntimeException(e);
-        }
         if (TestHelpers.isInLauncherProcess()) {
             Utilities.enableRunningInTestHarnessForTests();
             mLauncher.setSystemHealthSupplier(startTime -> TestCommandReceiver.callCommand(
                             TestCommandReceiver.GET_SYSTEM_HEALTH_MESSAGE, startTime.toString())
                     .getString("result"));
         }
-        mLauncher.enableDebugTracing();
-        // Avoid double-reporting of Launcher crashes.
-        mLauncher.setOnLauncherCrashed(() -> mLauncherPid = 0);
     }
 
-    @Rule
-    public ShellCommandRule mDisableHeadsUpNotification =
-            ShellCommandRule.disableHeadsUpNotification();
-
-    @Rule
-    public ScreenRecordRule mScreenRecordRule = new ScreenRecordRule();
-
-    @Rule
-    public SetFlagsRule mSetFlagsRule = new SetFlagsRule(DEVICE_DEFAULT);
-
-    @Rule
-    public ExtendedLongPressTimeoutRule mLongPressTimeoutRule = new ExtendedLongPressTimeoutRule();
-
-    @Rule
-    public LimitDevicesRule mlimitDevicesRule = new LimitDevicesRule();
-
+    /**
+     * @deprecated call {@link #performInitialization} instead
+     */
+    @Deprecated
     public static void initialize(AbstractLauncherUiTest test) throws Exception {
-        test.reinitializeLauncherData();
-        test.mDevice.pressHome();
-        test.waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
-        test.waitForState("Launcher internal state didn't switch to Home",
-                () -> LauncherState.NORMAL);
-        test.waitForResumed("Launcher internal state is still Background");
+        test.performInitialization();
+    }
+
+    @Override
+    protected void performInitialization() {
+        reinitializeLauncherData();
+        mDevice.pressHome();
         // Check that we switched to home.
-        test.mLauncher.getWorkspace();
-        AbstractLauncherUiTest.checkDetectedLeaks(test.mLauncher, true);
+        mLauncher.getWorkspace();
+
+        waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
+        waitForState("Launcher internal state didn't switch to Home",
+                () -> LauncherState.NORMAL);
+        waitForResumed("Launcher internal state is still Background");
+
+        checkDetectedLeaks(mLauncher, true);
     }
 
-    protected void clearPackageData(String pkg) throws IOException, InterruptedException {
-        assertTrue("pm clear command failed",
-                mDevice.executeShellCommand("pm clear " + pkg)
-                .contains("Success"));
-        assertTrue("pm wait-for-handler command failed",
-                mDevice.executeShellCommand("pm wait-for-handler")
-                .contains("Success"));
-    }
-
+    @Override
     protected TestRule getRulesInsideActivityMonitor() {
         final ViewCaptureRule viewCaptureRule = new ViewCaptureRule(
                 Launcher.ACTIVITY_TRACKER::getCreatedContext);
         final RuleChain inner = RuleChain
-                .outerRule(new PortraitLandscapeRunner<LAUNCHER_TYPE>(this))
+                .outerRule(new PortraitLandscapeRunner<>(this))
                 .around(new FailureWatcher(mLauncher, viewCaptureRule::getViewCaptureData))
                 // .around(viewCaptureRule) // b/315482167
                 .around(new TestIsolationRule(mLauncher, true));
@@ -265,184 +117,6 @@
                 : inner;
     }
 
-    @Rule
-    public TestRule mOrderSensitiveRules = RuleChain
-            .outerRule(new SamplerRule())
-            .around(new TestStabilityRule())
-            .around(getRulesInsideActivityMonitor());
-
-    public UiDevice getDevice() {
-        return mDevice;
-    }
-
-    @Before
-    public void setUp() throws Exception {
-        mLauncher.onTestStart();
-
-        final String launcherPackageName = mDevice.getLauncherPackageName();
-        try {
-            final Context context = InstrumentationRegistry.getContext();
-            final PackageManager pm = context.getPackageManager();
-            final PackageInfo launcherPackage = pm.getPackageInfo(launcherPackageName, 0);
-
-            if (!launcherPackage.versionName.equals("BuildFromAndroidStudio")) {
-                Assert.assertEquals("Launcher version doesn't match tests version",
-                        pm.getPackageInfo(context.getPackageName(), 0).getLongVersionCode(),
-                        launcherPackage.getLongVersionCode());
-            }
-        } catch (PackageManager.NameNotFoundException e) {
-            throw new RuntimeException(e);
-        }
-
-        mLauncherPid = 0;
-
-        mTargetContext = InstrumentationRegistry.getTargetContext();
-        mTargetPackage = mTargetContext.getPackageName();
-        mLauncherPid = mLauncher.getPid();
-
-        UserManager userManager = mTargetContext.getSystemService(UserManager.class);
-        if (userManager != null) {
-            for (UserHandle userHandle : userManager.getUserProfiles()) {
-                if (!userHandle.isSystem()) {
-                    mDevice.executeShellCommand(
-                            "pm remove-user --wait " + userHandle.getIdentifier());
-                }
-            }
-        }
-
-        onTestStart();
-
-        initialize(this);
-    }
-
-    private long getAvailableMemory() {
-        mActivityManager.getMemoryInfo(mMemoryInfo);
-
-        return Math.divideExact(mMemoryInfo.availMem,  BYTES_PER_MEGABYTE);
-    }
-
-    @Before
-    public void saveMemoryBefore() {
-        mMemoryBefore = getAvailableMemory();
-    }
-
-    @After
-    public void logMemoryAfter() {
-        long memoryAfter = getAvailableMemory();
-
-        Log.d(TAG, "Available memory: before=" + mMemoryBefore
-                + "MB, after=" + memoryAfter
-                + "MB, delta=" + (memoryAfter - mMemoryBefore) + "MB");
-    }
-
-    /** Method that should be called when a test starts. */
-    public static void onTestStart() {
-        waitForSetupWizardDismissal();
-
-        if (TestStabilityRule.isPresubmit()) {
-            aggressivelyUnlockSysUi();
-        } else {
-            verifyKeyguardInvisible();
-        }
-    }
-
-    private static boolean hasSystemUiObject(String resId) {
-        return getUiDevice().hasObject(
-                By.res(SYSTEMUI_PACKAGE, resId));
-    }
-
-    @NonNull
-    private static UiDevice getUiDevice() {
-        return UiDevice.getInstance(getInstrumentation());
-    }
-
-    private static void aggressivelyUnlockSysUi() {
-        final UiDevice device = getUiDevice();
-        for (int i = 0; i < 10 && hasSystemUiObject("keyguard_status_view"); ++i) {
-            Log.d(TAG, "Before attempting to unlock the phone");
-            try {
-                device.executeShellCommand("input keyevent 82");
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
-            device.waitForIdle();
-        }
-        Assert.assertTrue("Keyguard still visible",
-                TestHelpers.wait(
-                        Until.gone(By.res(SYSTEMUI_PACKAGE, "keyguard_status_view")), 60000));
-        Log.d(TAG, "Keyguard is not visible");
-    }
-
-    /** Waits for setup wizard to go away. */
-    private static void waitForSetupWizardDismissal() {
-        if (sFirstTimeWaitingForWizard) {
-            try {
-                getUiDevice().executeShellCommand(
-                        "am force-stop com.google.android.setupwizard");
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
-        }
-
-        final boolean wizardDismissed = TestHelpers.wait(
-                Until.gone(By.pkg("com.google.android.setupwizard").depth(0)),
-                sFirstTimeWaitingForWizard ? 120000 : 0);
-        sFirstTimeWaitingForWizard = false;
-        Assert.assertTrue("Setup wizard is still visible", wizardDismissed);
-    }
-
-    /** Asserts that keyguard is not visible */
-    public static void verifyKeyguardInvisible() {
-        final boolean keyguardAlreadyVisible = sSeenKeyguard;
-
-        sSeenKeyguard = sSeenKeyguard
-                || !TestHelpers.wait(
-                Until.gone(By.res(SYSTEMUI_PACKAGE, "keyguard_status_view")), 60000);
-
-        Assert.assertFalse(
-                "Keyguard is visible, which is likely caused by a crash in SysUI, seeing keyguard"
-                        + " for the first time = "
-                        + !keyguardAlreadyVisible,
-                sSeenKeyguard);
-    }
-
-    @After
-    public void resetFreezeRecentTaskList() {
-        try {
-            mDevice.executeShellCommand("wm reset-freeze-recent-tasks");
-        } catch (IOException e) {
-            Log.e(TAG, "Failed to reset fozen recent tasks list", e);
-        }
-    }
-
-    @After
-    public void verifyLauncherState() {
-        try {
-            // Limits UI tests affecting tests running after them.
-            mDevice.pressHome();
-            mLauncher.waitForLauncherInitialized();
-            if (mLauncherPid != 0) {
-                assertEquals("Launcher crashed, pid mismatch:",
-                        mLauncherPid, mLauncher.getPid().intValue());
-            }
-        } finally {
-            mLauncher.onTestFinish();
-        }
-    }
-
-    protected void reinitializeLauncherData() {
-        reinitializeLauncherData(false);
-    }
-
-    protected void reinitializeLauncherData(boolean clearWorkspace) {
-        if (clearWorkspace) {
-            mLauncher.clearLauncherData();
-        } else {
-            mLauncher.reinitializeLauncherData();
-        }
-        mLauncher.waitForLauncherInitialized();
-    }
-
     /**
      * Runs the callback on the UI thread and returns the result.
      */
@@ -668,45 +342,4 @@
 
     protected void onLauncherActivityClose(LAUNCHER_TYPE launcher) {
     }
-
-    protected HomeAppIcon createShortcutInCenterIfNotExist(String name) {
-        Point dimension = mLauncher.getWorkspace().getIconGridDimensions();
-        return createShortcutIfNotExist(name, dimension.x / 2, dimension.y / 2);
-    }
-
-    protected HomeAppIcon createShortcutIfNotExist(String name, Point cellPosition) {
-        return createShortcutIfNotExist(name, cellPosition.x, cellPosition.y);
-    }
-
-    protected HomeAppIcon createShortcutIfNotExist(String name, int cellX, int cellY) {
-        HomeAppIcon homeAppIcon = mLauncher.getWorkspace().tryGetWorkspaceAppIcon(name);
-        Log.d(ICON_MISSING, "homeAppIcon: " + homeAppIcon + " name: " + name +
-                " cell: " + cellX + ", " + cellY);
-        if (homeAppIcon == null) {
-            HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-            allApps.freeze();
-            try {
-                allApps.getAppIcon(name).dragToWorkspace(cellX, cellY);
-            } finally {
-                allApps.unfreeze();
-            }
-            homeAppIcon = mLauncher.getWorkspace().getWorkspaceAppIcon(name);
-        }
-        return homeAppIcon;
-    }
-
-    protected void commitTransactionAndLoadHome(FavoriteItemsTransaction transaction) {
-        transaction.commit();
-
-        // Launch the home activity
-        UiDevice.getInstance(getInstrumentation()).pressHome();
-        mLauncher.waitForLauncherInitialized();
-    }
-
-    /** Clears all recent tasks */
-    protected void clearAllRecentTasks() {
-        if (!mLauncher.getRecentTasks().isEmpty()) {
-            mLauncher.goHome().switchToOverview().dismissAllTasks();
-        }
-    }
 }
diff --git a/tests/src/com/android/launcher3/ui/BaseLauncherTaplTest.java b/tests/src/com/android/launcher3/ui/BaseLauncherTaplTest.java
new file mode 100644
index 0000000..8449853
--- /dev/null
+++ b/tests/src/com/android/launcher3/ui/BaseLauncherTaplTest.java
@@ -0,0 +1,529 @@
+/*
+ * 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.ui;
+
+import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
+
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.testing.shared.TestProtocol.ICON_MISSING;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import android.app.ActivityManager;
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.ActivityInfo;
+import android.content.pm.PackageInfo;
+import android.content.pm.PackageManager;
+import android.graphics.Point;
+import android.os.Debug;
+import android.os.RemoteException;
+import android.os.UserHandle;
+import android.os.UserManager;
+import android.platform.test.flag.junit.SetFlagsRule;
+import android.platform.test.rule.LimitDevicesRule;
+import android.util.Log;
+
+import androidx.annotation.NonNull;
+import androidx.test.InstrumentationRegistry;
+import androidx.test.uiautomator.By;
+import androidx.test.uiautomator.BySelector;
+import androidx.test.uiautomator.UiDevice;
+import androidx.test.uiautomator.Until;
+
+import com.android.launcher3.celllayout.FavoriteItemsTransaction;
+import com.android.launcher3.tapl.HomeAllApps;
+import com.android.launcher3.tapl.HomeAppIcon;
+import com.android.launcher3.tapl.LauncherInstrumentation;
+import com.android.launcher3.tapl.TestHelpers;
+import com.android.launcher3.util.TestUtil;
+import com.android.launcher3.util.Wait;
+import com.android.launcher3.util.rule.ExtendedLongPressTimeoutRule;
+import com.android.launcher3.util.rule.FailureWatcher;
+import com.android.launcher3.util.rule.SamplerRule;
+import com.android.launcher3.util.rule.ScreenRecordRule;
+import com.android.launcher3.util.rule.ShellCommandRule;
+import com.android.launcher3.util.rule.TestIsolationRule;
+import com.android.launcher3.util.rule.TestStabilityRule;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.rules.RuleChain;
+import org.junit.rules.TestRule;
+
+import java.io.IOException;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Base class for all TAPL tests in Launcher providing various utility methods.
+ */
+public abstract class BaseLauncherTaplTest {
+
+    public static final long DEFAULT_ACTIVITY_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
+    public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 10;
+
+    public static final long DEFAULT_UI_TIMEOUT = TestUtil.DEFAULT_UI_TIMEOUT;
+    private static final String TAG = "BaseLauncherTaplTest";
+
+    private static final long BYTES_PER_MEGABYTE = 1 << 20;
+
+    private static boolean sDumpWasGenerated = false;
+    private static boolean sActivityLeakReported = false;
+    private static boolean sSeenKeyguard = false;
+    private static boolean sFirstTimeWaitingForWizard = true;
+
+    private static final String SYSTEMUI_PACKAGE = "com.android.systemui";
+
+    protected final UiDevice mDevice = getUiDevice();
+    protected final LauncherInstrumentation mLauncher = createLauncherInstrumentation();
+
+    @NonNull
+    public static LauncherInstrumentation createLauncherInstrumentation() {
+        waitForSetupWizardDismissal(); // precondition for creating LauncherInstrumentation
+        return new LauncherInstrumentation(true);
+    }
+
+    protected Context mTargetContext;
+    protected String mTargetPackage;
+    private int mLauncherPid;
+
+    private final ActivityManager.MemoryInfo mMemoryInfo = new ActivityManager.MemoryInfo();
+    private final ActivityManager mActivityManager;
+    private long mMemoryBefore;
+
+    /** Detects activity leaks and throws an exception if a leak is found. */
+    public static void checkDetectedLeaks(LauncherInstrumentation launcher) {
+        checkDetectedLeaks(launcher, false);
+    }
+
+    /** Detects activity leaks and throws an exception if a leak is found. */
+    public static void checkDetectedLeaks(LauncherInstrumentation launcher,
+            boolean requireOneActiveActivityUnused) {
+        if (TestStabilityRule.isPresubmit()) return; // b/313501215
+
+        final boolean requireOneActiveActivity =
+                false; // workaround for leaks when there is an unexpected Recents activity
+
+        if (sActivityLeakReported) return;
+
+        // Check whether activity leak detector has found leaked activities.
+        Wait.atMost(() -> getActivityLeakErrorMessage(launcher, requireOneActiveActivity),
+                () -> {
+                    launcher.forceGc();
+                    return MAIN_EXECUTOR.submit(
+                            () -> launcher.noLeakedActivities(requireOneActiveActivity)).get();
+                }, launcher, DEFAULT_UI_TIMEOUT);
+    }
+
+    public static String getAppPackageName() {
+        return getInstrumentation().getContext().getPackageName();
+    }
+
+    private static String getActivityLeakErrorMessage(LauncherInstrumentation launcher,
+            boolean requireOneActiveActivity) {
+        sActivityLeakReported = true;
+        return "Activity leak detector has found leaked activities, requirining 1 activity: "
+                + requireOneActiveActivity + "; "
+                + dumpHprofData(launcher, false, requireOneActiveActivity) + ".";
+    }
+
+    private static String dumpHprofData(LauncherInstrumentation launcher, boolean intentionalLeak,
+            boolean requireOneActiveActivity) {
+        if (intentionalLeak) return "intentional leak; not generating dump";
+
+        String result;
+        if (sDumpWasGenerated) {
+            result = "dump has already been generated by another test";
+        } else {
+            try {
+                final String fileName =
+                        getInstrumentation().getTargetContext().getFilesDir().getPath()
+                                + "/ActivityLeakHeapDump.hprof";
+                if (TestHelpers.isInLauncherProcess()) {
+                    Debug.dumpHprofData(fileName);
+                } else {
+                    final UiDevice device = getUiDevice();
+                    device.executeShellCommand(
+                            "am dumpheap " + device.getLauncherPackageName() + " " + fileName);
+                }
+                Log.d(TAG, "Saved leak dump, the leak is still present: "
+                        + !launcher.noLeakedActivities(requireOneActiveActivity));
+                sDumpWasGenerated = true;
+                result = "saved memory dump as an artifact";
+            } catch (Throwable e) {
+                Log.e(TAG, "dumpHprofData failed", e);
+                result = "failed to save memory dump";
+            }
+        }
+        return result + ". Full list of activities: " + launcher.getRootedActivitiesList();
+    }
+
+    protected BaseLauncherTaplTest() {
+        mActivityManager = InstrumentationRegistry.getContext()
+                .getSystemService(ActivityManager.class);
+        mLauncher.enableCheckEventsForSuccessfulGestures();
+        mLauncher.setAnomalyChecker(BaseLauncherTaplTest::verifyKeyguardInvisible);
+        try {
+            mDevice.setOrientationNatural();
+        } catch (RemoteException e) {
+            throw new RuntimeException(e);
+        }
+        mLauncher.enableDebugTracing();
+        // Avoid double-reporting of Launcher crashes.
+        mLauncher.setOnLauncherCrashed(() -> mLauncherPid = 0);
+    }
+
+    @Rule
+    public ShellCommandRule mDisableHeadsUpNotification =
+            ShellCommandRule.disableHeadsUpNotification();
+
+    @Rule
+    public ScreenRecordRule mScreenRecordRule = new ScreenRecordRule();
+
+    @Rule
+    public SetFlagsRule mSetFlagsRule = new SetFlagsRule(DEVICE_DEFAULT);
+
+    @Rule
+    public ExtendedLongPressTimeoutRule mLongPressTimeoutRule = new ExtendedLongPressTimeoutRule();
+
+    @Rule
+    public LimitDevicesRule mlimitDevicesRule = new LimitDevicesRule();
+
+    protected void performInitialization() {
+        reinitializeLauncherData();
+        mDevice.pressHome();
+        // Check that we switched to home.
+        mLauncher.getWorkspace();
+        checkDetectedLeaks(mLauncher, true);
+    }
+
+    protected void clearPackageData(String pkg) throws IOException, InterruptedException {
+        assertTrue("pm clear command failed",
+                mDevice.executeShellCommand("pm clear " + pkg)
+                        .contains("Success"));
+        assertTrue("pm wait-for-handler command failed",
+                mDevice.executeShellCommand("pm wait-for-handler")
+                        .contains("Success"));
+    }
+
+    protected TestRule getRulesInsideActivityMonitor() {
+        final RuleChain inner = RuleChain
+                .outerRule(new FailureWatcher(mLauncher, null))
+                .around(new TestIsolationRule(mLauncher, true));
+        return TestHelpers.isInLauncherProcess()
+                ? RuleChain.outerRule(ShellCommandRule.setDefaultLauncher()).around(inner)
+                : inner;
+    }
+
+    @Rule
+    public TestRule mOrderSensitiveRules = RuleChain
+            .outerRule(new SamplerRule())
+            .around(new TestStabilityRule())
+            .around(getRulesInsideActivityMonitor());
+
+    public UiDevice getDevice() {
+        return mDevice;
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        mLauncher.onTestStart();
+
+        final String launcherPackageName = mDevice.getLauncherPackageName();
+        try {
+            final Context context = InstrumentationRegistry.getContext();
+            final PackageManager pm = context.getPackageManager();
+            final PackageInfo launcherPackage = pm.getPackageInfo(launcherPackageName, 0);
+
+            if (!launcherPackage.versionName.equals("BuildFromAndroidStudio")) {
+                Assert.assertEquals("Launcher version doesn't match tests version",
+                        pm.getPackageInfo(context.getPackageName(), 0).getLongVersionCode(),
+                        launcherPackage.getLongVersionCode());
+            }
+        } catch (PackageManager.NameNotFoundException e) {
+            throw new RuntimeException(e);
+        }
+
+        mLauncherPid = 0;
+
+        mTargetContext = InstrumentationRegistry.getTargetContext();
+        mTargetPackage = mTargetContext.getPackageName();
+        mLauncherPid = mLauncher.getPid();
+
+        UserManager userManager = mTargetContext.getSystemService(UserManager.class);
+        if (userManager != null) {
+            for (UserHandle userHandle : userManager.getUserProfiles()) {
+                if (!userHandle.isSystem()) {
+                    mDevice.executeShellCommand(
+                            "pm remove-user --wait " + userHandle.getIdentifier());
+                }
+            }
+        }
+
+        onTestStart();
+        performInitialization();
+    }
+
+    private long getAvailableMemory() {
+        mActivityManager.getMemoryInfo(mMemoryInfo);
+
+        return Math.divideExact(mMemoryInfo.availMem,  BYTES_PER_MEGABYTE);
+    }
+
+    @Before
+    public void saveMemoryBefore() {
+        mMemoryBefore = getAvailableMemory();
+    }
+
+    @After
+    public void logMemoryAfter() {
+        long memoryAfter = getAvailableMemory();
+
+        Log.d(TAG, "Available memory: before=" + mMemoryBefore
+                + "MB, after=" + memoryAfter
+                + "MB, delta=" + (memoryAfter - mMemoryBefore) + "MB");
+    }
+
+    /** Method that should be called when a test starts. */
+    public static void onTestStart() {
+        waitForSetupWizardDismissal();
+
+        if (TestStabilityRule.isPresubmit()) {
+            aggressivelyUnlockSysUi();
+        } else {
+            verifyKeyguardInvisible();
+        }
+    }
+
+    private static boolean hasSystemUiObject(String resId) {
+        return getUiDevice().hasObject(
+                By.res(SYSTEMUI_PACKAGE, resId));
+    }
+
+    @NonNull
+    private static UiDevice getUiDevice() {
+        return UiDevice.getInstance(getInstrumentation());
+    }
+
+    private static void aggressivelyUnlockSysUi() {
+        final UiDevice device = getUiDevice();
+        for (int i = 0; i < 10 && hasSystemUiObject("keyguard_status_view"); ++i) {
+            Log.d(TAG, "Before attempting to unlock the phone");
+            try {
+                device.executeShellCommand("input keyevent 82");
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            device.waitForIdle();
+        }
+        Assert.assertTrue("Keyguard still visible",
+                TestHelpers.wait(
+                        Until.gone(By.res(SYSTEMUI_PACKAGE, "keyguard_status_view")), 60000));
+        Log.d(TAG, "Keyguard is not visible");
+    }
+
+    /** Waits for setup wizard to go away. */
+    private static void waitForSetupWizardDismissal() {
+        if (sFirstTimeWaitingForWizard) {
+            try {
+                getUiDevice().executeShellCommand(
+                        "am force-stop com.google.android.setupwizard");
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        final boolean wizardDismissed = TestHelpers.wait(
+                Until.gone(By.pkg("com.google.android.setupwizard").depth(0)),
+                sFirstTimeWaitingForWizard ? 120000 : 0);
+        sFirstTimeWaitingForWizard = false;
+        Assert.assertTrue("Setup wizard is still visible", wizardDismissed);
+    }
+
+    /** Asserts that keyguard is not visible */
+    public static void verifyKeyguardInvisible() {
+        final boolean keyguardAlreadyVisible = sSeenKeyguard;
+
+        sSeenKeyguard = sSeenKeyguard
+                || !TestHelpers.wait(
+                Until.gone(By.res(SYSTEMUI_PACKAGE, "keyguard_status_view")), 60000);
+
+        Assert.assertFalse(
+                "Keyguard is visible, which is likely caused by a crash in SysUI, seeing keyguard"
+                        + " for the first time = "
+                        + !keyguardAlreadyVisible,
+                sSeenKeyguard);
+    }
+
+    @After
+    public void resetFreezeRecentTaskList() {
+        try {
+            mDevice.executeShellCommand("wm reset-freeze-recent-tasks");
+        } catch (IOException e) {
+            Log.e(TAG, "Failed to reset fozen recent tasks list", e);
+        }
+    }
+
+    @After
+    public void verifyLauncherState() {
+        try {
+            // Limits UI tests affecting tests running after them.
+            mDevice.pressHome();
+            mLauncher.waitForLauncherInitialized();
+            if (mLauncherPid != 0) {
+                assertEquals("Launcher crashed, pid mismatch:",
+                        mLauncherPid, mLauncher.getPid().intValue());
+            }
+        } finally {
+            mLauncher.onTestFinish();
+        }
+    }
+
+    protected void reinitializeLauncherData() {
+        reinitializeLauncherData(false);
+    }
+
+    protected void reinitializeLauncherData(boolean clearWorkspace) {
+        if (clearWorkspace) {
+            mLauncher.clearLauncherData();
+        } else {
+            mLauncher.reinitializeLauncherData();
+        }
+        mLauncher.waitForLauncherInitialized();
+    }
+
+    public static void startAppFast(String packageName) {
+        startIntent(
+                getInstrumentation().getContext().getPackageManager().getLaunchIntentForPackage(
+                        packageName),
+                By.pkg(packageName).depth(0),
+                true /* newTask */);
+    }
+
+    public static void startTestActivity(String activityName, String activityLabel) {
+        final String packageName = getAppPackageName();
+        final Intent intent = getInstrumentation().getContext().getPackageManager()
+                        .getLaunchIntentForPackage(packageName);
+        intent.setComponent(new ComponentName(packageName,
+                "com.android.launcher3.tests." + activityName));
+        startIntent(intent, By.pkg(packageName).text(activityLabel),
+                false /* newTask */);
+    }
+
+    public static void startTestActivity(int activityNumber) {
+        startTestActivity("Activity" + activityNumber, "TestActivity" + activityNumber);
+    }
+
+    public static void startImeTestActivity() {
+        final String packageName = getAppPackageName();
+        final Intent intent = getInstrumentation().getContext().getPackageManager()
+                        .getLaunchIntentForPackage(packageName);
+        intent.setComponent(new ComponentName(packageName,
+                "com.android.launcher3.testcomponent.ImeTestActivity"));
+        startIntent(intent, By.pkg(packageName).text("ImeTestActivity"),
+                false /* newTask */);
+    }
+
+    /** Starts ExcludeFromRecentsTestActivity, which has excludeFromRecents="true". */
+    public static void startExcludeFromRecentsTestActivity() {
+        final String packageName = getAppPackageName();
+        final Intent intent = getInstrumentation().getContext().getPackageManager()
+                .getLaunchIntentForPackage(packageName);
+        intent.setComponent(new ComponentName(packageName,
+                "com.android.launcher3.testcomponent.ExcludeFromRecentsTestActivity"));
+        startIntent(intent, By.pkg(packageName).text("ExcludeFromRecentsTestActivity"),
+                false /* newTask */);
+    }
+
+    private static void startIntent(Intent intent, BySelector selector, boolean newTask) {
+        intent.addCategory(Intent.CATEGORY_LAUNCHER);
+        if (newTask) {
+            intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK);
+        } else {
+            intent.addFlags(
+                    Intent.FLAG_ACTIVITY_MULTIPLE_TASK | Intent.FLAG_ACTIVITY_NEW_DOCUMENT);
+        }
+        getInstrumentation().getTargetContext().startActivity(intent);
+        assertTrue("App didn't start: " + selector,
+                TestHelpers.wait(Until.hasObject(selector), DEFAULT_UI_TIMEOUT));
+
+        // Wait for the Launcher to stop.
+        final LauncherInstrumentation launcherInstrumentation = new LauncherInstrumentation();
+        Wait.atMost("Launcher activity didn't stop",
+                () -> !launcherInstrumentation.isLauncherActivityStarted(),
+                launcherInstrumentation, DEFAULT_ACTIVITY_TIMEOUT);
+    }
+
+    public static ActivityInfo resolveSystemAppInfo(String category) {
+        return getInstrumentation().getContext().getPackageManager().resolveActivity(
+                new Intent(Intent.ACTION_MAIN).addCategory(category),
+                PackageManager.MATCH_SYSTEM_ONLY)
+                .activityInfo;
+    }
+
+
+    public static String resolveSystemApp(String category) {
+        return resolveSystemAppInfo(category).packageName;
+    }
+
+    protected HomeAppIcon createShortcutInCenterIfNotExist(String name) {
+        Point dimension = mLauncher.getWorkspace().getIconGridDimensions();
+        return createShortcutIfNotExist(name, dimension.x / 2, dimension.y / 2);
+    }
+
+    protected HomeAppIcon createShortcutIfNotExist(String name, Point cellPosition) {
+        return createShortcutIfNotExist(name, cellPosition.x, cellPosition.y);
+    }
+
+    protected HomeAppIcon createShortcutIfNotExist(String name, int cellX, int cellY) {
+        HomeAppIcon homeAppIcon = mLauncher.getWorkspace().tryGetWorkspaceAppIcon(name);
+        Log.d(ICON_MISSING, "homeAppIcon: " + homeAppIcon + " name: " + name
+                + " cell: " + cellX + ", " + cellY);
+        if (homeAppIcon == null) {
+            HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
+            allApps.freeze();
+            try {
+                allApps.getAppIcon(name).dragToWorkspace(cellX, cellY);
+            } finally {
+                allApps.unfreeze();
+            }
+            homeAppIcon = mLauncher.getWorkspace().getWorkspaceAppIcon(name);
+        }
+        return homeAppIcon;
+    }
+
+    protected void commitTransactionAndLoadHome(FavoriteItemsTransaction transaction) {
+        transaction.commit();
+
+        // Launch the home activity
+        UiDevice.getInstance(getInstrumentation()).pressHome();
+        mLauncher.waitForLauncherInitialized();
+    }
+
+    /** Clears all recent tasks */
+    protected void clearAllRecentTasks() {
+        if (!mLauncher.getRecentTasks().isEmpty()) {
+            mLauncher.goHome().switchToOverview().dismissAllTasks();
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
index 7bdc040..3b85309 100644
--- a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
+++ b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
@@ -12,7 +12,7 @@
 
 import com.android.app.viewcapture.data.ExportedData;
 import com.android.launcher3.tapl.LauncherInstrumentation;
-import com.android.launcher3.ui.AbstractLauncherUiTest;
+import com.android.launcher3.ui.BaseLauncherTaplTest;
 
 import org.junit.rules.TestWatcher;
 import org.junit.runner.Description;
@@ -57,7 +57,7 @@
     @Override
     protected void succeeded(Description description) {
         super.succeeded(description);
-        AbstractLauncherUiTest.checkDetectedLeaks(mLauncher);
+        BaseLauncherTaplTest.checkDetectedLeaks(mLauncher);
     }
 
     @Override
diff --git a/tests/tapl/com/android/launcher3/tapl/KeyboardQuickSwitch.java b/tests/tapl/com/android/launcher3/tapl/KeyboardQuickSwitch.java
index 7ff55fe..7cb2614 100644
--- a/tests/tapl/com/android/launcher3/tapl/KeyboardQuickSwitch.java
+++ b/tests/tapl/com/android/launcher3/tapl/KeyboardQuickSwitch.java
@@ -163,6 +163,32 @@
     }
 
     /**
+     * Dismisses the Keyboard Quick Switch view by going home. After the Keyboard Quick Switch view
+     * gets hidden, it unpresses ALT key, which is generally used to keep the view visible.
+     */
+    public Workspace dismissByGoingHome() {
+        try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
+                "verifying keyboard quick switch view is shown")) {
+            mLauncher.waitForLauncherObject(KEYBOARD_QUICK_SWITCH_RES_ID);
+        }
+
+        mLauncher.goHome();
+
+        try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
+                "waiting for keyboard quick switch dismissal");
+             LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
+            mLauncher.waitUntilLauncherObjectGone(KEYBOARD_QUICK_SWITCH_RES_ID);
+        }
+
+        try (LauncherInstrumentation.Closable c2 = mLauncher.addContextLayer(
+                "get workspace after releasing ALT key")) {
+            mLauncher.expectEvent(TestProtocol.SEQUENCE_MAIN, EVENT_HOME_ALT_LEFT_UP);
+            mLauncher.unpressKeyCode(KeyEvent.KEYCODE_ALT_LEFT, 0);
+            return mLauncher.getWorkspace();
+        }
+    }
+
+    /**
      * Launches the currently-focused app task.
      * <p>
      * This method should only be used if the focused task is for a recent running app, otherwise