Merge "Add a flag for the context menu to pin/unpin apps on taskbar." into main
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 782a705..6367a01 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -361,7 +361,11 @@
     <dimen name="taskbar_running_app_indicator_width">12dp</dimen>
     <dimen name="taskbar_running_app_indicator_top_margin">4dp</dimen>
     <dimen name="taskbar_minimized_app_indicator_width">6dp</dimen>
-    <dimen name="taskbar_overflow_button_preview_stroke">2dp</dimen>
+    <dimen name="taskbar_overflow_item_icon_size_default">22dp</dimen>
+    <dimen name="taskbar_overflow_item_icon_size_scaled_down">15dp</dimen>
+    <dimen name="taskbar_overflow_item_icon_stroke_width_default">2dp</dimen>
+    <dimen name="taskbar_overflow_leave_behind_size_default">18dp</dimen>
+    <dimen name="taskbar_overflow_leave_behind_size_scaled_down">15dp</dimen>
 
     <!-- Transient taskbar -->
     <dimen name="transient_taskbar_padding">12dp</dimen>
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/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 82acc0c..d7e5c61 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -418,7 +418,7 @@
     /** Called when the visibility of the bubble bar changed. */
     public void bubbleBarVisibilityChanged(boolean isVisible) {
         mControllers.uiController.adjustHotseatForBubbleBar(isVisible);
-        mControllers.taskbarViewController.resetIconAlignmentController();
+        mControllers.taskbarViewController.adjustTaskbarForBubbleBar();
     }
 
     public void init(@NonNull TaskbarSharedState sharedState) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarOverflowView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarOverflowView.java
index 126e9bb..712478e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarOverflowView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarOverflowView.java
@@ -16,17 +16,28 @@
 
 package com.android.launcher3.taskbar;
 
+import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
+import android.animation.AnimatorSet;
+import android.animation.ObjectAnimator;
+import android.animation.ValueAnimator;
 import android.content.Context;
+import android.graphics.BlendMode;
+import android.graphics.BlendModeColorFilter;
 import android.graphics.Canvas;
 import android.graphics.Paint;
 import android.graphics.drawable.Drawable;
 import android.util.AttributeSet;
+import android.util.FloatProperty;
+import android.util.IntProperty;
 import android.view.LayoutInflater;
 import android.view.ViewGroup;
 import android.widget.FrameLayout;
 
 import androidx.annotation.NonNull;
+import androidx.core.graphics.ColorUtils;
 
+import com.android.app.animation.Interpolators;
 import com.android.launcher3.R;
 import com.android.launcher3.Reorderable;
 import com.android.launcher3.Utilities;
@@ -45,8 +56,104 @@
  * each other in counter clockwise manner (icons of tasks partially overlapping with each other).
  */
 public class TaskbarOverflowView extends FrameLayout implements Reorderable {
+    private static final int ALPHA_TRANSPARENT = 0;
+    private static final int ALPHA_OPAQUE = 255;
+    private static final long ANIMATION_DURATION_APPS_TO_LEAVE_BEHIND = 300L;
+    private static final long ANIMATION_DURATION_LEAVE_BEHIND_TO_APPS = 500L;
+    private static final long ANIMATION_SET_DURATION = 1000L;
+    private static final long ITEM_ICON_CENTER_OFFSET_ANIMATION_DURATION = 500L;
+    private static final long ITEM_ICON_COLOR_FILTER_OPACITY_ANIMATION_DURATION = 600L;
+    private static final long ITEM_ICON_SIZE_ANIMATION_DURATION = 500L;
+    private static final long ITEM_ICON_STROKE_WIDTH_ANIMATION_DURATION = 500L;
+    private static final long LEAVE_BEHIND_ANIMATIONS_DELAY = 500L;
+    private static final long LEAVE_BEHIND_OPACITY_ANIMATION_DURATION = 100L;
+    private static final long LEAVE_BEHIND_SIZE_ANIMATION_DURATION = 500L;
     private static final int MAX_ITEMS_IN_PREVIEW = 4;
 
+    private static final FloatProperty<TaskbarOverflowView> ITEM_ICON_CENTER_OFFSET =
+            new FloatProperty<>("itemIconCenterOffset") {
+                @Override
+                public Float get(TaskbarOverflowView view) {
+                    return view.mItemIconCenterOffset;
+                }
+
+                @Override
+                public void setValue(TaskbarOverflowView view, float value) {
+                    view.mItemIconCenterOffset = value;
+                    view.invalidate();
+                }
+            };
+
+    private static final IntProperty<TaskbarOverflowView> ITEM_ICON_COLOR_FILTER_OPACITY =
+            new IntProperty<>("itemIconColorFilterOpacity") {
+                @Override
+                public Integer get(TaskbarOverflowView view) {
+                    return view.mItemIconColorFilterOpacity;
+                }
+
+                @Override
+                public void setValue(TaskbarOverflowView view, int value) {
+                    view.mItemIconColorFilterOpacity = value;
+                    view.invalidate();
+                }
+            };
+
+    private static final FloatProperty<TaskbarOverflowView> ITEM_ICON_SIZE =
+            new FloatProperty<>("itemIconSize") {
+                @Override
+                public Float get(TaskbarOverflowView view) {
+                    return view.mItemIconSize;
+                }
+
+                @Override
+                public void setValue(TaskbarOverflowView view, float value) {
+                    view.mItemIconSize = value;
+                    view.invalidate();
+                }
+            };
+
+    private static final FloatProperty<TaskbarOverflowView> ITEM_ICON_STROKE_WIDTH =
+            new FloatProperty<>("itemIconStrokeWidth") {
+                @Override
+                public Float get(TaskbarOverflowView view) {
+                    return view.mItemIconStrokeWidth;
+                }
+
+                @Override
+                public void setValue(TaskbarOverflowView view, float value) {
+                    view.mItemIconStrokeWidth = value;
+                    view.invalidate();
+                }
+            };
+
+    private static final IntProperty<TaskbarOverflowView> LEAVE_BEHIND_OPACITY =
+            new IntProperty<>("leaveBehindOpacity") {
+                @Override
+                public Integer get(TaskbarOverflowView view) {
+                    return view.mLeaveBehindOpacity;
+                }
+
+                @Override
+                public void setValue(TaskbarOverflowView view, int value) {
+                    view.mLeaveBehindOpacity = value;
+                    view.invalidate();
+                }
+            };
+
+    private static final FloatProperty<TaskbarOverflowView> LEAVE_BEHIND_SIZE =
+            new FloatProperty<>("leaveBehindSize") {
+                @Override
+                public Float get(TaskbarOverflowView view) {
+                    return view.mLeaveBehindSize;
+                }
+
+                @Override
+                public void setValue(TaskbarOverflowView view, float value) {
+                    view.mLeaveBehindSize = value;
+                    view.invalidate();
+                }
+            };
+
     private boolean mIsRtlLayout;
     private final List<Task> mItems = new ArrayList<Task>();
     private int mIconSize;
@@ -56,11 +163,24 @@
     private float mScaleForReorderBounce = 1f;
     private int mItemBackgroundColor;
     private int mLeaveBehindColor;
-    private float mItemPreviewStrokeWidth;
 
     // Active means the overflow icon has been pressed, which replaces the app icons with the
     // leave-behind circle and shows the KQS UI.
     private boolean mIsActive = false;
+    private ValueAnimator mStateTransitionAnimationWrapper;
+
+    private float mItemIconCenterOffsetDefault;
+    private float mItemIconCenterOffset;  // [0..mItemIconCenterOffsetDefault]
+    private int mItemIconColorFilterOpacity;  // [ALPHA_TRANSPARENT..ALPHA_OPAQUE]
+    private float mItemIconSizeDefault;
+    private float mItemIconSizeScaledDown;
+    private float mItemIconSize;  // [mItemIconSizeScaledDown..mItemIconSizeDefault]
+    private float mItemIconStrokeWidthDefault;
+    private float mItemIconStrokeWidth;  // [0..mItemIconStrokeWidthDefault]
+    private int mLeaveBehindOpacity;  // [ALPHA_TRANSPARENT..ALPHA_OPAQUE]
+    private float mLeaveBehindSizeScaledDown;
+    private float mLeaveBehindSizeDefault;
+    private float mLeaveBehindSize;  // [mLeaveBehindSizeScaledDown..mLeaveBehindSizeDefault]
 
     public TaskbarOverflowView(Context context, AttributeSet attrs) {
         super(context, attrs);
@@ -87,6 +207,12 @@
 
         icon.mIconSize = iconSize;
         icon.mPadding = padding;
+
+        final float radius = iconSize / 2f - padding;
+        final float size = radius + icon.mItemIconStrokeWidth;
+        icon.mItemIconCenterOffsetDefault = radius - size / 2 - icon.mItemIconStrokeWidth;
+        icon.mItemIconCenterOffset = icon.mItemIconCenterOffsetDefault;
+
         return icon;
     }
 
@@ -95,8 +221,22 @@
         mItemBackgroundPaint = new Paint(Paint.ANTI_ALIAS_FLAG);
         mItemBackgroundColor = getContext().getColor(R.color.taskbar_background);
         mLeaveBehindColor = Themes.getAttrColor(getContext(), android.R.attr.textColorTertiary);
-        mItemPreviewStrokeWidth = getResources().getDimension(
-                R.dimen.taskbar_overflow_button_preview_stroke);
+
+        mItemIconSizeDefault = getResources().getDimension(
+                R.dimen.taskbar_overflow_item_icon_size_default);
+        mItemIconSizeScaledDown = getResources().getDimension(
+                R.dimen.taskbar_overflow_item_icon_size_scaled_down);
+        mItemIconSize = mItemIconSizeDefault;
+
+        mItemIconStrokeWidthDefault = getResources().getDimension(
+                R.dimen.taskbar_overflow_item_icon_stroke_width_default);
+        mItemIconStrokeWidth = mItemIconStrokeWidthDefault;
+
+        mLeaveBehindSizeDefault = getResources().getDimension(
+                R.dimen.taskbar_overflow_leave_behind_size_default);
+        mLeaveBehindSizeScaledDown = getResources().getDimension(
+                R.dimen.taskbar_overflow_leave_behind_size_scaled_down);
+        mLeaveBehindSize = mLeaveBehindSizeScaledDown;
 
         setWillNotDraw(false);
     }
@@ -105,16 +245,14 @@
     protected void onDraw(@NonNull Canvas canvas) {
         super.onDraw(canvas);
 
-        if (mIsActive) {
-            drawLeaveBehindCircle(canvas);
-        } else {
-            drawAppIcons(canvas);
-        }
+        drawAppIcons(canvas);
+        drawLeaveBehindCircle(canvas);
     }
 
     private void drawAppIcons(@NonNull Canvas canvas) {
         mItemBackgroundPaint.setColor(mItemBackgroundColor);
         float radius = mIconSize / 2f - mPadding;
+        int adjustedItemIconSize = Math.round(mItemIconSize);
 
         int itemsToShow = Math.min(mItems.size(), MAX_ITEMS_IN_PREVIEW);
         for (int i = itemsToShow - 1; i >= 0; --i) {
@@ -123,36 +261,33 @@
                 continue;
             }
 
-            // Set the item icon size so two items fit within the overflow icon with stroke width
-            // included, and overlap of 4 stroke width sizes between base item preview items.
-            // 2 * strokeWidth + 2 * itemIconSize - 4 * strokeWidth = iconSize = 2 * radius.
-            float itemIconSize = radius + mItemPreviewStrokeWidth;
-            // Offset item icon from center so item icon stroke edge matches the parent icon edge.
-            float itemCenterOffset = radius - itemIconSize / 2 - mItemPreviewStrokeWidth;
-
-            float itemCenterX = getItemXOffset(itemCenterOffset, mIsRtlLayout, i, itemsToShow);
-            float itemCenterY = getItemYOffset(itemCenterOffset, i, itemsToShow);
+            float itemCenterX = getItemXOffset(mItemIconCenterOffset, mIsRtlLayout, i, itemsToShow);
+            float itemCenterY = getItemYOffset(mItemIconCenterOffset, i, itemsToShow);
 
             Drawable iconCopy = icon.getConstantState().newDrawable().mutate();
-            iconCopy.setBounds(0, 0, (int) itemIconSize, (int) itemIconSize);
+            iconCopy.setBounds(0, 0, adjustedItemIconSize, adjustedItemIconSize);
+            iconCopy.setColorFilter(new BlendModeColorFilter(
+                    ColorUtils.setAlphaComponent(mLeaveBehindColor, mItemIconColorFilterOpacity),
+                    BlendMode.SRC_ATOP));
 
             canvas.save();
-            float itemIconRadius = itemIconSize / 2;
+            float itemIconRadius = adjustedItemIconSize / 2f;
             canvas.translate(
                     mPadding + itemCenterX + radius - itemIconRadius,
                     mPadding + itemCenterY + radius - itemIconRadius);
             canvas.drawCircle(itemIconRadius, itemIconRadius,
-                    itemIconRadius + mItemPreviewStrokeWidth, mItemBackgroundPaint);
+                    itemIconRadius + mItemIconStrokeWidth, mItemBackgroundPaint);
             iconCopy.draw(canvas);
             canvas.restore();
         }
     }
 
     private void drawLeaveBehindCircle(@NonNull Canvas canvas) {
-        mItemBackgroundPaint.setColor(mLeaveBehindColor);
+        mItemBackgroundPaint.setColor(
+                ColorUtils.setAlphaComponent(mLeaveBehindColor, mLeaveBehindOpacity));
 
-        final var xyCenter = mIconSize / 2f;
-        canvas.drawCircle(xyCenter, xyCenter, mIconSize / 4f, mItemBackgroundPaint);
+        final float xyCenter = mIconSize / 2f;
+        canvas.drawCircle(xyCenter, xyCenter, mLeaveBehindSize / 2f, mItemBackgroundPaint);
     }
 
     /**
@@ -203,10 +338,98 @@
      * @param isActive The next state of the view.
      */
     public void setIsActive(boolean isActive) {
-        if (mIsActive != isActive) {
-            mIsActive = isActive;
-            invalidate();
+        if (mIsActive == isActive) {
+            return;
         }
+        mIsActive = isActive;
+
+        if (mStateTransitionAnimationWrapper != null
+                && mStateTransitionAnimationWrapper.isRunning()) {
+            mStateTransitionAnimationWrapper.reverse();
+            return;
+        }
+
+        final AnimatorSet stateTransitionAnimation = getStateTransitionAnimation();
+        mStateTransitionAnimationWrapper = ValueAnimator.ofFloat(0, 1f);
+        mStateTransitionAnimationWrapper.setDuration(mIsActive
+                ? ANIMATION_DURATION_APPS_TO_LEAVE_BEHIND
+                : ANIMATION_DURATION_LEAVE_BEHIND_TO_APPS);
+        mStateTransitionAnimationWrapper.setInterpolator(
+                mIsActive ? Interpolators.STANDARD : Interpolators.EMPHASIZED);
+        mStateTransitionAnimationWrapper.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                mStateTransitionAnimationWrapper = null;
+            }
+        });
+        mStateTransitionAnimationWrapper.addUpdateListener(
+                new ValueAnimator.AnimatorUpdateListener() {
+                    @Override
+                    public void onAnimationUpdate(ValueAnimator animator) {
+                        stateTransitionAnimation.setCurrentPlayTime(
+                                (long) (ANIMATION_SET_DURATION * animator.getAnimatedFraction()));
+                    }
+                });
+        mStateTransitionAnimationWrapper.start();
+    }
+
+    private AnimatorSet getStateTransitionAnimation() {
+        final AnimatorSet animation = new AnimatorSet();
+        animation.setInterpolator(Interpolators.LINEAR);
+        animation.playTogether(
+                buildAnimator(ITEM_ICON_CENTER_OFFSET, 0f, mItemIconCenterOffsetDefault,
+                        ITEM_ICON_CENTER_OFFSET_ANIMATION_DURATION, 0L,
+                        ITEM_ICON_CENTER_OFFSET_ANIMATION_DURATION),
+                buildAnimator(ITEM_ICON_COLOR_FILTER_OPACITY, ALPHA_OPAQUE, ALPHA_TRANSPARENT,
+                        ITEM_ICON_COLOR_FILTER_OPACITY_ANIMATION_DURATION, 0L,
+                        ANIMATION_SET_DURATION - ITEM_ICON_COLOR_FILTER_OPACITY_ANIMATION_DURATION),
+                buildAnimator(ITEM_ICON_SIZE, mItemIconSizeScaledDown, mItemIconSizeDefault,
+                        ITEM_ICON_SIZE_ANIMATION_DURATION, 0L,
+                        ITEM_ICON_SIZE_ANIMATION_DURATION),
+                buildAnimator(ITEM_ICON_STROKE_WIDTH, 0f, mItemIconStrokeWidthDefault,
+                        ITEM_ICON_STROKE_WIDTH_ANIMATION_DURATION, 0L,
+                        ITEM_ICON_STROKE_WIDTH_ANIMATION_DURATION),
+                buildAnimator(LEAVE_BEHIND_OPACITY, ALPHA_OPAQUE, ALPHA_TRANSPARENT,
+                        LEAVE_BEHIND_OPACITY_ANIMATION_DURATION, LEAVE_BEHIND_ANIMATIONS_DELAY,
+                        ANIMATION_SET_DURATION - LEAVE_BEHIND_ANIMATIONS_DELAY
+                                - LEAVE_BEHIND_OPACITY_ANIMATION_DURATION),
+                buildAnimator(LEAVE_BEHIND_SIZE, mLeaveBehindSizeDefault,
+                        mLeaveBehindSizeScaledDown, LEAVE_BEHIND_SIZE_ANIMATION_DURATION,
+                        LEAVE_BEHIND_ANIMATIONS_DELAY, 0L)
+        );
+        return animation;
+    }
+
+    private ObjectAnimator buildAnimator(IntProperty<TaskbarOverflowView> property,
+            int finalValueWhenAnimatingToLeaveBehind, int finalValueWhenAnimatingToAppIcons,
+            long duration, long delayWhenAnimatingToLeaveBehind,
+            long delayWhenAnimatingToAppIcons) {
+        final ObjectAnimator animator = ObjectAnimator.ofInt(this, property,
+                mIsActive ? finalValueWhenAnimatingToLeaveBehind
+                        : finalValueWhenAnimatingToAppIcons);
+        applyTiming(animator, duration, delayWhenAnimatingToLeaveBehind,
+                delayWhenAnimatingToAppIcons);
+        return animator;
+    }
+
+    private ObjectAnimator buildAnimator(FloatProperty<TaskbarOverflowView> property,
+            float finalValueWhenAnimatingToLeaveBehind, float finalValueWhenAnimatingToAppIcons,
+            long duration, long delayWhenAnimatingToLeaveBehind,
+            long delayWhenAnimatingToAppIcons) {
+        final ObjectAnimator animator = ObjectAnimator.ofFloat(this, property,
+                mIsActive ? finalValueWhenAnimatingToLeaveBehind
+                        : finalValueWhenAnimatingToAppIcons);
+        applyTiming(animator, duration, delayWhenAnimatingToLeaveBehind,
+                delayWhenAnimatingToAppIcons);
+        return animator;
+    }
+
+    private void applyTiming(ObjectAnimator animator, long duration,
+            long delayWhenAnimatingToLeaveBehind,
+            long delayWhenAnimatingToAppIcons) {
+        animator.setDuration(duration);
+        animator.setStartDelay(
+                mIsActive ? delayWhenAnimatingToLeaveBehind : delayWhenAnimatingToAppIcons);
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index 55bcb23..8816a6d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -120,7 +120,8 @@
 
     private boolean mShouldTryStartAlign;
 
-    private final int mMaxNumIcons;
+    private int mMaxNumIcons = 0;
+    private int mIdealNumIcons = 0;
 
     private final int mAllAppsButtonTranslationOffset;
 
@@ -188,8 +189,6 @@
 
         // TODO: Disable touch events on QSB otherwise it can crash.
         mQsb = LayoutInflater.from(context).inflate(R.layout.search_container_hotseat, this, false);
-
-        mMaxNumIcons = calculateMaxNumIcons();
     }
 
     /**
@@ -200,11 +199,15 @@
         int availableWidth = deviceProfile.widthPx;
         int defaultEdgeMargin =
                 (int) getResources().getDimension(deviceProfile.inv.inlineNavButtonsEndSpacing);
+        int spaceForBubbleBar =
+                Math.round(mControllerCallbacks.getBubbleBarMaxCollapsedWidthIfVisible());
 
         // Reserve space required for edge margins, or for navbar if shown. If task bar needs to be
         // center aligned with nav bar shown, reserve space on both sides.
-        availableWidth -= Math.max(defaultEdgeMargin, deviceProfile.hotseatBarEndOffset);
-        availableWidth -= Math.max(defaultEdgeMargin,
+        availableWidth -=
+                Math.max(defaultEdgeMargin + spaceForBubbleBar, deviceProfile.hotseatBarEndOffset);
+        availableWidth -= Math.max(
+                defaultEdgeMargin + (mShouldTryStartAlign ? 0 : spaceForBubbleBar),
                 mShouldTryStartAlign ? 0 : deviceProfile.hotseatBarEndOffset);
 
         // The space taken by an item icon used during layout.
@@ -231,6 +234,21 @@
         return Math.floorDiv(availableWidth, iconSize) + additionalIcons;
     }
 
+    /**
+     * Recalculates the max number of icons the taskbar view can show without entering overflow.
+     * Returns whether the max number of icons changed and the change affects the number of icons
+     * that should be shown in the taskbar.
+     */
+    boolean updateMaxNumIcons() {
+        if (!Flags.taskbarOverflow()) {
+            return false;
+        }
+        int oldMaxNumIcons = mMaxNumIcons;
+        mMaxNumIcons = calculateMaxNumIcons();
+        return oldMaxNumIcons != mMaxNumIcons
+                && (mIdealNumIcons > oldMaxNumIcons || mIdealNumIcons > mMaxNumIcons);
+    }
+
     @Override
     public void setVisibility(int visibility) {
         boolean changed = getVisibility() != visibility;
@@ -328,6 +346,10 @@
                 && mActivityContext.getTaskbarFeatureEvaluator().getSupportsPinningPopup()) {
             setOnTouchListener(mControllerCallbacks.getTaskbarTouchListener());
         }
+
+        if (Flags.taskbarOverflow()) {
+            mMaxNumIcons = calculateMaxNumIcons();
+        }
     }
 
     private void removeAndRecycle(View view) {
@@ -460,8 +482,9 @@
                 }
             }
 
-            overflowSize =
-                    nextViewIndex + numberOfSupportedRecents + nonTaskIconsToBeAdded - mMaxNumIcons;
+            mIdealNumIcons = nextViewIndex + numberOfSupportedRecents + nonTaskIconsToBeAdded;
+            overflowSize = mIdealNumIcons - mMaxNumIcons;
+
             if (overflowSize > 0 && mTaskbarOverflowView != null) {
                 addView(mTaskbarOverflowView, nextViewIndex++);
             } else if (mTaskbarOverflowView != null) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacks.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacks.java
index 834f92e..f65f307 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacks.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewCallbacks.java
@@ -137,6 +137,17 @@
         return null;
     }
 
+    /**
+     * Get the max bubble bar collapsed width for the current bubble bar visibility state. Used to
+     * reserve space for the bubble bar when transitioning taskbar view into overflow.
+     */
+    public float getBubbleBarMaxCollapsedWidthIfVisible() {
+        return mControllers.bubbleControllers
+                .filter(c -> !c.bubbleBarViewController.isHiddenForNoBubbles())
+                .map(c -> c.bubbleBarViewController.getCollapsedWidthWithMaxVisibleBubbles())
+                .orElse(0f);
+    }
+
     /** Returns true if bubble bar controllers present and enabled in persistent taskbar. */
     public boolean isBubbleBarEnabledInPersistentTaskbar() {
         return Flags.enableBubbleBarInPersistentTaskBar()
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
index 494c472..cebabff 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
@@ -779,9 +779,16 @@
         }
     }
 
-    /** Resets the icon alignment controller so that it can be recreated again later. */
-    void resetIconAlignmentController() {
+    /**
+     * Resets the icon alignment controller so that it can be recreated again later, and updates
+     * the list of icons shown in the taskbar if the bubble bar visibility changes the taskbar
+     * overflow state.
+     */
+    void adjustTaskbarForBubbleBar() {
         mIconAlignControllerLazy = null;
+        if (mTaskbarView.updateMaxNumIcons()) {
+            commitRunningAppsToUI();
+        }
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index 350f56f..c5c2d69 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -1290,10 +1290,14 @@
         // If there are more than 2 bubbles, the first 2 should be visible when collapsed,
         // excluding the overflow.
         return bubbleChildCount >= MAX_VISIBLE_BUBBLES_COLLAPSED
-                ? getScaledIconSize() + mIconOverlapAmount + horizontalPadding
+                ? getCollapsedWidthWithMaxVisibleBubbles()
                 : getScaledIconSize() + horizontalPadding;
     }
 
+    float getCollapsedWidthWithMaxVisibleBubbles()  {
+        return getScaledIconSize() + mIconOverlapAmount + 2 * mBubbleBarPadding;
+    }
+
     /** Returns the child count excluding the overflow if it's present. */
     int getBubbleChildCount() {
         return hasOverflow() ? getChildCount() - 1 : getChildCount();
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index 2dddb16..d842138 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -473,6 +473,13 @@
     }
 
     /**
+     * @return the max collapsed width for the bubble bar.
+     */
+    public float getCollapsedWidthWithMaxVisibleBubbles() {
+        return mBarView.getCollapsedWidthWithMaxVisibleBubbles();
+    }
+
+    /**
      * @return {@code true} if bubble bar is on the left edge of the screen, {@code false} if on
      * the right
      */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
index 7b20eea..908e97c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
@@ -21,6 +21,7 @@
 import android.view.ViewGroup
 import android.widget.FrameLayout
 import androidx.core.animation.ValueAnimator
+import com.android.app.animation.InterpolatorsAndroidX
 import com.android.launcher3.R
 import com.android.systemui.util.addListener
 
@@ -35,7 +36,8 @@
 ) {
 
     private companion object {
-        const val ANIMATION_DURATION_MS = 250L
+        const val EXPAND_ANIMATION_DURATION_MS = 400L
+        const val COLLAPSE_ANIMATION_DURATION_MS = 350L
     }
 
     private var flyout: BubbleBarFlyoutView? = null
@@ -86,9 +88,10 @@
     private fun showFlyout(animationType: AnimationType, endAction: () -> Unit) {
         val flyout = this.flyout ?: return
         val startValue = getCurrentAnimatedValueIfRunning() ?: 0f
-        val duration = (ANIMATION_DURATION_MS * (1f - startValue)).toLong()
+        val duration = (EXPAND_ANIMATION_DURATION_MS * (1f - startValue)).toLong()
         animator?.cancel()
         val animator = ValueAnimator.ofFloat(startValue, 1f).setDuration(duration)
+        animator.interpolator = InterpolatorsAndroidX.EMPHASIZED
         this.animator = animator
         when (animationType) {
             AnimationType.FADE ->
@@ -111,6 +114,7 @@
     fun updateFlyoutFullyExpanded(message: BubbleBarFlyoutMessage, onEnd: () -> Unit) {
         val flyout = flyout ?: return
         hideFlyout(AnimationType.FADE) {
+            callbacks.resetTopBoundary()
             flyout.updateData(message) { showFlyout(AnimationType.FADE, onEnd) }
         }
     }
@@ -152,9 +156,10 @@
     private fun hideFlyout(animationType: AnimationType, endAction: () -> Unit) {
         val flyout = this.flyout ?: return
         val startValue = getCurrentAnimatedValueIfRunning() ?: 1f
-        val duration = (ANIMATION_DURATION_MS * startValue).toLong()
+        val duration = (COLLAPSE_ANIMATION_DURATION_MS * startValue).toLong()
         animator?.cancel()
         val animator = ValueAnimator.ofFloat(startValue, 0f).setDuration(duration)
+        animator.interpolator = InterpolatorsAndroidX.EMPHASIZED
         this.animator = animator
         when (animationType) {
             AnimationType.FADE ->
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
index 418675c..f9f5a15 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
@@ -35,6 +35,7 @@
 import androidx.core.animation.ArgbEvaluator
 import com.android.launcher3.R
 import com.android.launcher3.popup.RoundedArrowDrawable
+import kotlin.math.min
 
 /** The flyout view used to notify the user of a new bubble notification. */
 class BubbleBarFlyoutView(
@@ -46,6 +47,8 @@
     private companion object {
         // the minimum progress of the expansion animation before the content starts fading in.
         const val MIN_EXPANSION_PROGRESS_FOR_CONTENT_ALPHA = 0.75f
+        // the rate multiple for the background color animation relative to the morph animation.
+        const val BACKGROUND_COLOR_CHANGE_RATE = 5
     }
 
     private val scheduler: FlyoutScheduler = scheduler ?: HandlerScheduler(this)
@@ -204,6 +207,8 @@
         minExpansionProgressForTriangle =
             positioner.distanceToRevealTriangle / translationToCollapsedPosition.y
 
+        backgroundPaint.color = collapsedColor
+
         // post the request to start the expand animation to the looper so the view can measure
         // itself
         scheduler.runAfterLayout(expandAnimation)
@@ -307,8 +312,16 @@
             height.toFloat() - triangleHeight + triangleOverlap,
         )
 
+        // transform the flyout color between the collapsed and expanded states. the color
+        // transformation completes at a faster rate (BACKGROUND_COLOR_CHANGE_RATE) than the
+        // expansion animation. this helps make the color change smooth.
         backgroundPaint.color =
-            ArgbEvaluator.getInstance().evaluate(expansionProgress, collapsedColor, backgroundColor)
+            ArgbEvaluator.getInstance()
+                .evaluate(
+                    min(expansionProgress * BACKGROUND_COLOR_CHANGE_RATE, 1f),
+                    collapsedColor,
+                    backgroundColor,
+                )
 
         canvas.save()
         canvas.translate(backgroundRectTx, backgroundRectTy)
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt
index 6632721..45f5568 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt
@@ -54,7 +54,9 @@
             if (field == state) return
             val transitionFromHome = field == BubbleLauncherState.HOME
             field = state
-            if (!bubbleBarViewController.hasBubbles()) {
+            val hasBubbles = bubbleBarViewController.hasBubbles()
+            bubbleBarViewController.onBubbleBarConfigurationChanged(hasBubbles)
+            if (!hasBubbles) {
                 // if there are no bubbles, there's nothing to show, so just return.
                 return
             }
@@ -65,7 +67,6 @@
                 // on home but in persistent taskbar elsewhere so the position is different.
                 animateBubbleBarY()
             }
-            bubbleBarViewController.onBubbleBarConfigurationChanged(/* animate= */ true)
         }
 
     override var isSysuiLocked: Boolean = false
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt
index 71303f8..e62c0d4 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt
@@ -87,7 +87,9 @@
         set(state) {
             if (field == state) return
             field = state
-            if (!bubbleBarViewController.hasBubbles()) {
+            val hasBubbles = bubbleBarViewController.hasBubbles()
+            bubbleBarViewController.onBubbleBarConfigurationChanged(hasBubbles)
+            if (!hasBubbles) {
                 // if there are no bubbles, there's nothing to show, so just return.
                 return
             }
@@ -103,7 +105,6 @@
             // Only stash if we're in an app, otherwise we're in home or overview where we should
             // be un-stashed
             updateStashedAndExpandedState(field == BubbleLauncherState.IN_APP, expand = false)
-            bubbleBarViewController.onBubbleBarConfigurationChanged(/* animate= */ true)
         }
 
     override var isSysuiLocked: Boolean = false
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/BaseContainerInterface.java b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
index a3953ca..2164bc2 100644
--- a/quickstep/src/com/android/quickstep/BaseContainerInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
@@ -378,6 +378,9 @@
     public static void getTaskDimension(Context context, DeviceProfile dp, PointF out) {
         out.x = dp.widthPx;
         out.y = dp.heightPx;
+        if (dp.isTablet && !DisplayController.isTransientTaskbar(context)) {
+            out.y -= dp.taskbarHeight;
+        }
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/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/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/util/SplitSelectStateController.java b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
index ea582c4..d35a36a 100644
--- a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
@@ -163,6 +163,8 @@
      */
     private Pair<InstanceId, com.android.launcher3.logging.InstanceId> mSessionInstanceIds;
 
+    private boolean mIsDestroyed = false;
+
     private final BackPressHandler mSplitBackHandler = new BackPressHandler() {
         @Override
         public boolean canHandleBack() {
@@ -199,6 +201,7 @@
 
     public void onDestroy() {
         mContainer = null;
+        mIsDestroyed = true;
         mActivityBackCallback = null;
         mAppPairsController.onDestroy();
         mSplitSelectDataHolder.onDestroy();
@@ -744,7 +747,9 @@
      */
     public void resetState() {
         mSplitSelectDataHolder.resetState();
-        mContainer.<RecentsView>getOverviewPanel().resetDesktopTaskFromSplitSelectState();
+        if (!mIsDestroyed) {
+            mContainer.<RecentsView>getOverviewPanel().resetDesktopTaskFromSplitSelectState();
+        }
         dispatchOnSplitSelectionExit();
         mRecentsAnimationRunning = false;
         mLaunchingTaskView = null;
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 3a4e328..9d74bfb 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1966,7 +1966,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 +2749,7 @@
             updateSizeAndPadding();
         }
 
-        showCurrentTask(mActiveGestureRunningTasks);
+        showCurrentTask(mActiveGestureRunningTasks, "onGestureAnimationStart");
         setEnableFreeScroll(false);
         setEnableDrawingLiveTile(false);
         setRunningTaskHidden(true);
@@ -2930,8 +2930,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;
         }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
index 48f3fc2..582ea54 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
@@ -1123,7 +1123,7 @@
             animator.animateBubbleInForStashed(updatedBubble, isExpanding = false)
 
             // the flyout should now reverse and expand
-            animatorTestRule.advanceTimeBy(100)
+            animatorTestRule.advanceTimeBy(400)
         }
 
         assertThat(flyoutView!!.findViewById<TextView>(R.id.bubble_flyout_text).text)
@@ -1362,21 +1362,21 @@
 
     private fun waitForFlyoutToShow() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            animatorTestRule.advanceTimeBy(250)
+            animatorTestRule.advanceTimeBy(400)
         }
         assertThat(flyoutView).isNotNull()
     }
 
     private fun waitForFlyoutToHide() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            animatorTestRule.advanceTimeBy(250)
+            animatorTestRule.advanceTimeBy(350)
         }
         assertThat(flyoutView).isNull()
     }
 
     private fun waitForFlyoutToFadeOutAndBackIn() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            animatorTestRule.advanceTimeBy(500)
+            animatorTestRule.advanceTimeBy(750)
         }
         assertThat(flyoutView).isNotNull()
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
index 2997ac9..103c769 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
@@ -50,6 +50,9 @@
     private var onLeft = true
     private var flyoutTy = 50f
 
+    private val showAnimationDuration = 400L
+    private val hideAnimationDuration = 350L
+
     @Before
     fun setUp() {
         flyoutContainer = FrameLayout(context)
@@ -118,7 +121,7 @@
             assertThat(flyoutController.hasFlyout()).isTrue()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             flyoutController.collapseFlyout {}
-            animatorTestRule.advanceTimeBy(300)
+            animatorTestRule.advanceTimeBy(hideAnimationDuration)
         }
         assertThat(flyoutContainer.childCount).isEqualTo(0)
         assertThat(flyoutController.hasFlyout()).isFalse()
@@ -135,7 +138,7 @@
         }
         InstrumentationRegistry.getInstrumentation().waitForIdleSync()
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            animatorTestRule.advanceTimeBy(300)
+            animatorTestRule.advanceTimeBy(showAnimationDuration)
         }
         assertThat(flyoutCallbacks.topBoundaryExtendedSpace).isEqualTo(50)
     }
@@ -148,7 +151,7 @@
         }
         InstrumentationRegistry.getInstrumentation().waitForIdleSync()
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            animatorTestRule.advanceTimeBy(300)
+            animatorTestRule.advanceTimeBy(showAnimationDuration)
         }
         assertThat(flyoutCallbacks.topBoundaryExtendedSpace).isEqualTo(0)
     }
@@ -159,7 +162,7 @@
             setupAndShowFlyout()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             flyoutController.collapseFlyout {}
-            animatorTestRule.advanceTimeBy(300)
+            animatorTestRule.advanceTimeBy(hideAnimationDuration)
         }
         assertThat(flyoutCallbacks.topBoundaryReset).isTrue()
     }
@@ -172,7 +175,7 @@
             val flyoutView = flyoutContainer.findViewById<View>(R.id.bubble_bar_flyout_view)
             assertThat(flyoutView.alpha).isEqualTo(1f)
             flyoutController.cancelFlyout {}
-            animatorTestRule.advanceTimeBy(300)
+            animatorTestRule.advanceTimeBy(hideAnimationDuration)
             assertThat(flyoutView.alpha).isEqualTo(0f)
         }
         assertThat(flyoutCallbacks.topBoundaryReset).isTrue()
@@ -185,7 +188,7 @@
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             val flyoutView = flyoutContainer.findViewById<View>(R.id.bubble_bar_flyout_view)
             assertThat(flyoutView.alpha).isEqualTo(1f)
-            animatorTestRule.advanceTimeBy(300)
+            animatorTestRule.advanceTimeBy(showAnimationDuration)
             flyoutView.performClick()
         }
         assertThat(flyoutCallbacks.flyoutClicked).isTrue()
@@ -221,7 +224,7 @@
     fun updateFlyoutFullyExpanded() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
             setupAndShowFlyout()
-            animatorTestRule.advanceTimeBy(300)
+            animatorTestRule.advanceTimeBy(showAnimationDuration)
         }
         assertThat(flyoutController.hasFlyout()).isTrue()
 
@@ -234,13 +237,13 @@
             flyoutController.updateFlyoutFullyExpanded(newFlyoutMessage) {}
 
             // advance the timer so that the fade out animation plays
-            animatorTestRule.advanceTimeBy(250)
+            animatorTestRule.advanceTimeBy(hideAnimationDuration)
             assertThat(flyout.alpha).isEqualTo(0)
             assertThat(flyout.findViewById<TextView>(R.id.bubble_flyout_text).text)
                 .isEqualTo("new message")
 
             // advance the timer so that the fade in animation plays
-            animatorTestRule.advanceTimeBy(250)
+            animatorTestRule.advanceTimeBy(showAnimationDuration)
             assertThat(flyout.alpha).isEqualTo(1)
         }
         assertThat(flyoutCallbacks.topBoundaryExtendedSpace).isEqualTo(50)
@@ -250,7 +253,7 @@
     fun updateFlyoutWhileCollapsing() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
             setupAndShowFlyout()
-            animatorTestRule.advanceTimeBy(300)
+            animatorTestRule.advanceTimeBy(showAnimationDuration)
         }
         assertThat(flyoutController.hasFlyout()).isTrue()
 
@@ -265,9 +268,10 @@
             var flyoutReversed = false
             flyoutController.updateFlyoutWhileCollapsing(newFlyoutMessage) { flyoutReversed = true }
 
-            // the collapse animation ran for 125ms when it was updated, so reversing it should only
-            // run for the same amount of time
-            animatorTestRule.advanceTimeBy(125)
+            // the collapse and expand animations use an emphasized interpolator, so the reverse
+            // path does not take the same time. advance the timer the by full duration of the show
+            // animation to ensure it completes
+            animatorTestRule.advanceTimeBy(showAnimationDuration)
             val flyout = flyoutContainer.findViewById<View>(R.id.bubble_bar_flyout_view)
             assertThat(flyout.alpha).isEqualTo(1)
             assertThat(flyout.findViewById<TextView>(R.id.bubble_flyout_text).text)
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt
index b7ee6c4..f795ab1 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt
@@ -86,6 +86,20 @@
     }
 
     @Test
+    fun updateLauncherState_noBubbles_controllerNotified() {
+        // Given bubble bar has  no bubbles
+        whenever(bubbleBarViewController.hasBubbles()).thenReturn(false)
+
+        // When switch to home screen
+        getInstrumentation().runOnMainSync {
+            persistentTaskBarStashController.launcherState = BubbleLauncherState.HOME
+        }
+
+        // Then bubble bar view controller is notified
+        verify(bubbleBarViewController).onBubbleBarConfigurationChanged(/* animate= */ false)
+    }
+
+    @Test
     fun setBubblesShowingOnHomeUpdatedToFalse_barPositionYUpdated_controllersNotified() {
         // Given bubble bar is on home and has bubbles
         whenever(bubbleBarViewController.hasBubbles()).thenReturn(false)
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashControllerTest.kt
index 64416dd..1bbd12a 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashControllerTest.kt
@@ -119,6 +119,20 @@
     }
 
     @Test
+    fun updateLauncherState_noBubbles_controllerNotified() {
+        // Given bubble bar has  no bubbles
+        whenever(bubbleBarViewController.hasBubbles()).thenReturn(false)
+
+        // When switch to home screen
+        getInstrumentation().runOnMainSync {
+            mTransientBubbleStashController.launcherState = BubbleLauncherState.HOME
+        }
+
+        // Then bubble bar view controller is notified
+        verify(bubbleBarViewController).onBubbleBarConfigurationChanged(/* animate= */ false)
+    }
+
+    @Test
     fun setBubblesShowingOnHomeUpdatedToTrue_barPositionYUpdated_controllersNotified() {
         // Given bubble bar is on home and has bubbles
         whenever(bubbleBarViewController.hasBubbles()).thenReturn(true)
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
index dc5223c..6b95f8d 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
@@ -80,8 +80,6 @@
 
     protected final Context mContext =
             InstrumentationRegistry.getInstrumentation().getTargetContext();
-    protected final RecentsAnimationDeviceState mRecentsAnimationDeviceState =
-            new RecentsAnimationDeviceState(mContext, true);
     protected final InputConsumerController mInputConsumerController =
             InputConsumerController.getRecentsAnimationInputConsumer();
     protected final ActivityManager.RunningTaskInfo mRunningTaskInfo =
@@ -114,6 +112,7 @@
             new Bundle());
 
     protected TaskAnimationManager mTaskAnimationManager;
+    protected RecentsAnimationDeviceState mRecentsAnimationDeviceState;
 
     @Mock protected CONTAINER_INTERFACE mActivityInterface;
     @Mock protected ContextInitListener<?> mContextInitListener;
@@ -176,6 +175,12 @@
         }).when(recentsContainer).runOnBindToTouchInteractionService(any());
     }
 
+    @Before
+    public void setUpRecentsAnimationDeviceState() {
+        runOnMainSync(() ->
+                mRecentsAnimationDeviceState = new RecentsAnimationDeviceState(mContext, true));
+    }
+
     @Test
     public void testInitWhenReady_registersActivityInitListener() {
         String reasonString = "because i said so";
@@ -306,8 +311,6 @@
     }
 
     private void onRecentsAnimationStart(SWIPE_HANDLER absSwipeUpHandler) {
-        when(mActivityInterface.getOverviewWindowBounds(any(), any())).thenReturn(new Rect());
-
         runOnMainSync(() -> absSwipeUpHandler.onRecentsAnimationStart(
                 mRecentsAnimationController, mRecentsAnimationTargets));
     }
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/layout/widgets_list_expand_button.xml b/res/layout/widgets_list_expand_button.xml
index 17c19ac..ff2d777 100644
--- a/res/layout/widgets_list_expand_button.xml
+++ b/res/layout/widgets_list_expand_button.xml
@@ -15,6 +15,7 @@
   -->
 
 <Button xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/widget_list_expand_button"
     style="@style/Button.Rounded.Colored"
     android:layout_width="wrap_content"
     android:layout_height="wrap_content"
diff --git a/res/layout/work_mode_fab.xml b/res/layout/work_mode_fab.xml
index 46f2d8a..e2f0e09 100644
--- a/res/layout/work_mode_fab.xml
+++ b/res/layout/work_mode_fab.xml
@@ -17,6 +17,7 @@
     android:id="@+id/work_mode_toggle"
     android:layout_height="@dimen/work_fab_height"
     android:layout_width="wrap_content"
+    android:elevation="@dimen/work_fab_elevation"
     android:minHeight="@dimen/work_fab_height"
     android:gravity="center_vertical"
     android:background="@drawable/work_mode_fab_background"
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/res/values/dimens.xml b/res/values/dimens.xml
index c69778a..61d99d7 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -155,6 +155,7 @@
     <!-- Floating action button inside work tab to toggle work profile -->
     <dimen name="work_fab_height">56dp</dimen>
     <dimen name="work_fab_radius">16dp</dimen>
+    <dimen name="work_fab_elevation">6dp</dimen>
     <dimen name="work_fab_icon_size">24dp</dimen>
     <dimen name="work_fab_icon_vertical_margin">16dp</dimen>
     <dimen name="work_fab_icon_start_margin_expanded">4dp</dimen>
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 34cf56b..ef5c88a 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -20,6 +20,7 @@
 import static android.graphics.fonts.FontStyle.FONT_WEIGHT_NORMAL;
 import static android.text.Layout.Alignment.ALIGN_NORMAL;
 
+import static com.android.launcher3.Flags.enableContrastTiles;
 import static com.android.launcher3.Flags.enableCursorHoverStates;
 import static com.android.launcher3.graphics.PreloadIconDrawable.newPendingIcon;
 import static com.android.launcher3.icons.BitmapInfo.FLAG_NO_BADGE;
@@ -39,6 +40,7 @@
 import android.graphics.Color;
 import android.graphics.Paint;
 import android.graphics.Rect;
+import android.graphics.RectF;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.icu.text.MessageFormat;
@@ -722,6 +724,29 @@
         }
     }
 
+    /** Draws a background behind the App Title label when required. **/
+    public void drawAppContrastTile(Canvas canvas) {
+        RectF appTitleBounds;
+        Paint.FontMetrics fm = getPaint().getFontMetrics();
+        Rect tmpRect = new Rect();
+        getDrawingRect(tmpRect);
+
+        if (mIcon == null) {
+            appTitleBounds = new RectF(0, 0, tmpRect.right,
+                    (int) Math.ceil(fm.bottom - fm.top));
+        } else {
+            Rect iconBounds = new Rect();
+            getIconBounds(iconBounds);
+            int textStart = iconBounds.bottom + getCompoundDrawablePadding();
+            appTitleBounds = new RectF(tmpRect.left, textStart, tmpRect.right,
+                    textStart + (int) Math.ceil(fm.bottom - fm.top));
+        }
+
+        canvas.drawRoundRect(appTitleBounds, appTitleBounds.height() / 2,
+                appTitleBounds.height() / 2,
+                PillColorProvider.getInstance(getContext()).getAppTitlePillPaint());
+    }
+
     /** Draws a line under the app icon if this is representing a running app in Desktop Mode. */
     protected void drawRunningAppIndicatorIfNecessary(Canvas canvas) {
         if (mRunningAppState == RunningAppState.NOT_RUNNING || mDisplay != DISPLAY_TASKBAR) {
@@ -909,7 +934,9 @@
 
     @Override
     public void setTextColor(ColorStateList colors) {
-        mTextColor = colors.getDefaultColor();
+        mTextColor = shouldDrawAppContrastTile() ? PillColorProvider.getInstance(
+                getContext()).getAppTitleTextPaint().getColor()
+                : colors.getDefaultColor();
         mTextColorStateList = colors;
         if (Float.compare(mTextAlpha, 1) == 0) {
             super.setTextColor(colors);
@@ -926,6 +953,15 @@
                 && info.container != LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION);
     }
 
+    /**
+     * Whether or not an App title contrast tile should be drawn for this element.
+     **/
+    public boolean shouldDrawAppContrastTile() {
+        return mDisplay == DISPLAY_WORKSPACE && shouldTextBeVisible()
+                && PillColorProvider.getInstance(getContext()).isMatchaEnabled()
+                && enableContrastTiles();
+    }
+
     public void setTextVisibility(boolean visible) {
         setTextAlpha(visible ? 1 : 0);
     }
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/Launcher.java b/src/com/android/launcher3/Launcher.java
index 6145077..305941e 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -537,6 +537,7 @@
 
         mPopupDataProvider = new PopupDataProvider(this::updateNotificationDots);
         mWidgetPickerDataProvider = new WidgetPickerDataProvider();
+        PillColorProvider.getInstance(mWorkspace.getContext()).registerObserver();
 
         boolean internalStateHandled = ACTIVITY_TRACKER.handleCreate(this);
         if (internalStateHandled) {
@@ -1813,6 +1814,7 @@
         // changes while launcher is still loading.
         getRootView().getViewTreeObserver().removeOnPreDrawListener(mOnInitialBindListener);
         mOverlayManager.onActivityDestroyed();
+        PillColorProvider.getInstance(mWorkspace.getContext()).unregisterObserver();
     }
 
     public LauncherAccessibilityDelegate getAccessibilityDelegate() {
diff --git a/src/com/android/launcher3/PillColorPorovider.kt b/src/com/android/launcher3/PillColorPorovider.kt
new file mode 100644
index 0000000..347c5d6
--- /dev/null
+++ b/src/com/android/launcher3/PillColorPorovider.kt
@@ -0,0 +1,85 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3
+
+import android.content.Context
+import android.database.ContentObserver
+import android.graphics.Paint
+import android.net.Uri
+import android.provider.Settings
+import com.android.launcher3.util.Executors.ORDERED_BG_EXECUTOR
+
+class PillColorProvider private constructor(c: Context) {
+    private val context = c.applicationContext
+
+    private val matchaUri by lazy { Settings.Secure.getUriFor(MATCHA_SETTING) }
+    var appTitlePillPaint = Paint()
+        private set
+
+    var appTitleTextPaint = Paint()
+        private set
+
+    private var isMatchaEnabledInternal = 0
+
+    var isMatchaEnabled = isMatchaEnabledInternal != 0
+
+    private val pillColorObserver =
+        object : ContentObserver(ORDERED_BG_EXECUTOR.handler) {
+            override fun onChange(selfChange: Boolean, uri: Uri?) {
+                if (uri == matchaUri) {
+                    isMatchaEnabledInternal =
+                        Settings.Secure.getInt(context.contentResolver, MATCHA_SETTING, 0)
+                    isMatchaEnabled = isMatchaEnabledInternal != 0
+                }
+            }
+        }
+
+    fun registerObserver() {
+        context.contentResolver.registerContentObserver(matchaUri, false, pillColorObserver)
+        setup()
+    }
+
+    fun unregisterObserver() {
+        context.contentResolver.unregisterContentObserver(pillColorObserver)
+    }
+
+    fun setup() {
+        appTitlePillPaint.color =
+            context.resources.getColor(
+                R.color.material_color_surface_container_lowest,
+                context.theme,
+            )
+        appTitleTextPaint.color =
+            context.resources.getColor(R.color.material_color_on_surface, context.theme)
+        isMatchaEnabledInternal = Settings.Secure.getInt(context.contentResolver, MATCHA_SETTING, 0)
+        isMatchaEnabled = isMatchaEnabledInternal != 0
+    }
+
+    companion object {
+        private var INSTANCE: PillColorProvider? = null
+        private const val MATCHA_SETTING = "matcha_enable"
+
+        // TODO: Replace with a Dagger injection that is a singleton.
+        @JvmStatic
+        fun getInstance(context: Context): PillColorProvider {
+            if (INSTANCE == null) {
+                INSTANCE = PillColorProvider(context)
+            }
+            return INSTANCE!!
+        }
+    }
+}
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/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/views/DoubleShadowBubbleTextView.java b/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java
index ef66ffe..392d9a7 100644
--- a/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java
+++ b/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java
@@ -102,6 +102,9 @@
 
     @Override
     public void onDraw(Canvas canvas) {
+        if (shouldDrawAppContrastTile()) {
+            drawAppContrastTile(canvas);
+        }
         // If text is transparent or shadow alpha is 0, don't draw any shadow
         if (skipDoubleShadow()) {
             super.onDraw(canvas);
diff --git a/src/com/android/launcher3/views/StickyHeaderLayout.java b/src/com/android/launcher3/views/StickyHeaderLayout.java
index 090251f..4142e1f 100644
--- a/src/com/android/launcher3/views/StickyHeaderLayout.java
+++ b/src/com/android/launcher3/views/StickyHeaderLayout.java
@@ -120,7 +120,19 @@
     }
 
     private float getCurrentScroll() {
-        return mScrollOffset + (mCurrentEmptySpaceView == null ? 0 : mCurrentEmptySpaceView.getY());
+        float scroll;
+        if (mCurrentRecyclerView.getVisibility() != VISIBLE) {
+            // When no list is displayed, assume no scroll.
+            scroll = 0f;
+        } else if (mCurrentEmptySpaceView != null) {
+            // Otherwise use empty space view as reference to position.
+            scroll = mCurrentEmptySpaceView.getY();
+        } else {
+            // If there is no empty space view, but the list is visible, we are scrolled away
+            // completely, so assume all non-sticky children should also be scrolled away.
+            scroll = -mHeaderHeight;
+        }
+        return mScrollOffset + scroll;
     }
 
     @Override
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 2f64ab1..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();
@@ -605,9 +608,12 @@
             mAdapters.get(AdapterHolder.SEARCH).mWidgetsRecyclerView.setVisibility(GONE);
             mAdapters.get(getCurrentAdapterHolderType()).mWidgetsRecyclerView.setVisibility(
                     VISIBLE);
-            // Visibility of recommended widgets, recycler views and headers are handled in methods
-            // below.
-            post(this::onRecommendedWidgetsBound);
+            if (mRecommendedWidgetsCount > 0) {
+                // Display recommendations immediately, if present, so that other parts of sticky
+                // header (e.g. personal / work tabs) don't flash in interim.
+                mWidgetRecommendationsContainer.setVisibility(VISIBLE);
+            }
+            // Visibility of recycler views and headers are handled in methods below.
             onWidgetsBound();
         }
     }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index 3c67538..74a9a5c 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -103,7 +103,7 @@
                             .equals(mWidgetsContentVisiblePackageUserKey);
     @Nullable private Predicate<WidgetsListBaseEntry> mFilter = null;
     @Nullable private RecyclerView mRecyclerView;
-    @Nullable private PackageUserKey mPendingClickHeader;
+    @Nullable private PackageUserKey mHeaderPositionToMaintain;
     @Px private int mMaxHorizontalSpan;
 
     private boolean mShowOnlyDefaultList = true;
@@ -215,7 +215,7 @@
         // Get the current top of the header with the matching key before adjusting the visible
         // entries.
         OptionalInt previousPositionForPackageUserKey =
-                getPositionForPackageUserKey(mPendingClickHeader);
+                getPositionForPackageUserKey(mHeaderPositionToMaintain);
         OptionalInt topForPackageUserKey =
                 getOffsetForPosition(previousPositionForPackageUserKey);
 
@@ -247,13 +247,15 @@
         mVisibleEntries.addAll(newVisibleEntries);
         diffResult.dispatchUpdatesTo(this);
 
-        if (mPendingClickHeader != null) {
+        if (mHeaderPositionToMaintain != null && mRecyclerView != null) {
             // Get the position for the clicked header after adjusting the visible entries. The
             // position may have changed if another header had previously been expanded.
             OptionalInt positionForPackageUserKey =
-                    getPositionForPackageUserKey(mPendingClickHeader);
-            scrollToPositionAndMaintainOffset(positionForPackageUserKey, topForPackageUserKey);
-            mPendingClickHeader = null;
+                    getPositionForPackageUserKey(mHeaderPositionToMaintain);
+            // Post scroll updates to be applied after diff updates.
+            mRecyclerView.post(() -> scrollToPositionAndMaintainOffset(positionForPackageUserKey,
+                    topForPackageUserKey));
+            mHeaderPositionToMaintain = null;
         }
     }
 
@@ -384,7 +386,7 @@
 
         // Store the header that was clicked so that its position will be maintained the next time
         // we update the entries.
-        mPendingClickHeader = packageUserKey;
+        mHeaderPositionToMaintain = packageUserKey;
 
         updateVisibleEntries();
 
@@ -470,6 +472,16 @@
      */
     public void useExpandedList() {
         mShowOnlyDefaultList = false;
+        if (mWidgetsContentVisiblePackageUserKey != null) {
+            // Maintain selected header for the next update that expands the list.
+            mHeaderPositionToMaintain = mWidgetsContentVisiblePackageUserKey;
+        } else if (mVisibleEntries.size() > 2) {
+            // Maintain last visible header shown above expand button since there was no selected
+            // header.
+            mHeaderPositionToMaintain = PackageUserKey.fromPackageItemInfo(
+                    mVisibleEntries.get(mVisibleEntries.size() - 2).mPkgItem);
+        }
+
     }
 
     /** Comparator for sorting WidgetListRowEntry based on package title. */
diff --git a/tests/src/com/android/launcher3/dragging/TaplDragTest.java b/tests/src/com/android/launcher3/dragging/TaplDragTest.java
index 59e1f99..e2f9feb9a 100644
--- a/tests/src/com/android/launcher3/dragging/TaplDragTest.java
+++ b/tests/src/com/android/launcher3/dragging/TaplDragTest.java
@@ -64,7 +64,6 @@
     @Test
     @PortraitLandscape
     @PlatinumTest(focusArea = "launcher")
-    @ScreenRecordRule.ScreenRecord // b/353600888
     public void testDragToFolder() {
         // TODO: add the use case to drag an icon to an existing folder. Currently it either fails
         // on tablets or phones due to difference in resolution.
@@ -97,7 +96,6 @@
      * icon left.
      */
     @Test
-    @ScreenRecordRule.ScreenRecord // b/353600888
     public void testDragOutOfFolder() {
         final HomeAppIcon playStoreIcon = createShortcutIfNotExist(STORE_APP_NAME, 0, 1);
         final HomeAppIcon photosIcon = createShortcutInCenterIfNotExist(PHOTOS_APP_NAME);
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index a273648..2b1fddc 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -407,6 +407,15 @@
     }
 
     @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.
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
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 08c5552..fac73d3 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1728,6 +1728,27 @@
         scrollDownByDistance(container, distance, appsListBottomPadding);
     }
 
+    /** Scrolls up by given distance within the container. */
+    void scrollUpByDistance(UiObject2 container, int distance) {
+        scrollUpByDistance(container, distance, 0);
+    }
+
+    /** Scrolls up by given distance within the container considering the given bottom padding. */
+    void scrollUpByDistance(UiObject2 container, int distance, int bottomPadding) {
+        final Rect containerRect = getVisibleBounds(container);
+        final int bottomGestureMarginInContainer = getBottomGestureMarginInContainer(container);
+        scroll(
+                container,
+                Direction.UP,
+                new Rect(
+                        0,
+                        containerRect.height() - bottomGestureMarginInContainer - distance,
+                        0,
+                        bottomGestureMarginInContainer + bottomPadding),
+                /* steps= */ 10,
+                /* slowDown= */ true);
+    }
+
     void scrollDownByDistance(UiObject2 container, int distance) {
         scrollDownByDistance(container, distance, 0);
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 3097d9c..ac2748e 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -19,6 +19,7 @@
 import static com.android.launcher3.tapl.LauncherInstrumentation.WAIT_TIME_MS;
 import static com.android.launcher3.tapl.LauncherInstrumentation.log;
 
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.graphics.Rect;
 
@@ -31,6 +32,7 @@
 import com.android.launcher3.testing.shared.TestProtocol;
 
 import java.util.Collection;
+import java.util.List;
 
 /**
  * All widgets container.
@@ -128,8 +130,10 @@
             final UiObject2 searchBar = findSearchBar();
             final int searchBarHeight = searchBar.getVisibleBounds().height();
             final UiObject2 fullWidgetsPicker = verifyActiveContainer();
-            mLauncher.assertTrue("Widgets container didn't become scrollable",
-                    fullWidgetsPicker.wait(Until.scrollable(true), WAIT_TIME_MS));
+
+            // Widget picker may not be scrollable if there are few items. Instead of waiting on
+            // picker being scrollable, we wait on widget headers to be available.
+            waitForWidgetListItems(fullWidgetsPicker);
 
             final UiObject2 widgetsContainer =
                     findTestAppWidgetsTableContainer(testAppWidgetPackage);
@@ -176,6 +180,13 @@
         }
     }
 
+    private void waitForWidgetListItems(UiObject2 fullWidgetsPicker) {
+        List<UiObject2> headers = fullWidgetsPicker.wait(Until.findObjects(
+                By.res(mLauncher.getLauncherPackageName(), "widgets_list_header")), WAIT_TIME_MS);
+        mLauncher.assertTrue("Widgets list is not available",
+                headers != null && !headers.isEmpty());
+    }
+
     private UiObject2 findSearchBar() {
         final BySelector searchBarContainerSelector = By.res(mLauncher.getLauncherPackageName(),
                 "search_and_recommendations_container");
@@ -199,19 +210,38 @@
                 "container");
 
         String packageName =  mLauncher.getContext().getPackageName();
+        String packageNameToFind =
+                (testAppWidgetPackage == null || testAppWidgetPackage.isEmpty()) ? packageName
+                        : testAppWidgetPackage;
+
         final BySelector targetAppSelector = By
                 .clazz("android.widget.TextView")
-                .text((testAppWidgetPackage == null || testAppWidgetPackage.isEmpty())
-                                ? packageName
-                                : testAppWidgetPackage);
+                .text(packageNameToFind);
+        final BySelector expandListButtonSelector =
+                By.res(mLauncher.getLauncherPackageName(), "widget_list_expand_button");
         final BySelector widgetsContainerSelector = By.res(mLauncher.getLauncherPackageName(),
                 "widgets_table");
 
         boolean hasHeaderExpanded = false;
+        // List was expanded by clicking "Show all" button.
+        boolean hasListExpanded = false;
+
         int scrollDistance = 0;
         for (int i = 0; i < SCROLL_ATTEMPTS; i++) {
             UiObject2 widgetPicker = mLauncher.waitForLauncherObject(widgetPickerSelector);
             UiObject2 widgetListView = verifyActiveContainer();
+
+            // Press "Show all" button if it exists. Otherwise, keep scrolling to
+            // find the header or show all button.
+            UiObject2 expandListButton =
+                    mLauncher.findObjectInContainer(widgetListView, expandListButtonSelector);
+            if (expandListButton != null) {
+                expandListButton.click();
+                hasListExpanded = true;
+                i = -1;
+                continue;
+            }
+
             UiObject2 header = mLauncher.waitForObjectInContainer(widgetListView,
                     headerSelector);
             // If a header is barely visible in the bottom edge of the screen, its height could be
@@ -222,6 +252,17 @@
             // Look for a header that has the test app name.
             UiObject2 headerTitle = mLauncher.findObjectInContainer(widgetListView,
                     targetAppSelector);
+
+            final UiObject2 searchBar = findSearchBar();
+            // If header's title is under or above search bar, let's not process the header yet,
+            // scroll a bit more to bring the header into visible area.
+            if (headerTitle != null
+                    && headerTitle.getVisibleCenter().y <= searchBar.getVisibleCenter().y) {
+                log("Test app's header is behind the searchbar, scrolling up");
+                mLauncher.scrollUpByDistance(widgetListView, scrollDistance);
+                continue;
+            }
+
             if (headerTitle != null) {
                 // If we find the header and it has not been expanded, let's click it to see the
                 // widgets list. Note that we wait until the header is out of the gesture region at
@@ -258,11 +299,24 @@
                     widgetPicker,
                     widgetsContainerSelector);
 
-            mLauncher.scrollDownByDistance(hasHeaderExpanded && rightPane != null
-                    ? rightPane
-                    : widgetListView, scrollDistance);
+            if (hasListExpanded && packageNameToFind.compareToIgnoreCase(
+                    getFirstHeaderTitle(widgetListView)) < 0) {
+                mLauncher.scrollUpByDistance(hasHeaderExpanded && rightPane != null
+                        ? rightPane
+                        : widgetListView, scrollDistance);
+            } else {
+                mLauncher.scrollDownByDistance(hasHeaderExpanded && rightPane != null
+                        ? rightPane
+                        : widgetListView, scrollDistance);
+            }
         }
 
         return null;
     }
+
+    @NonNull
+    private String getFirstHeaderTitle(UiObject2 widgetListView) {
+        UiObject2 firstHeader = mLauncher.getObjectsInContainer(widgetListView, "app_title").get(0);
+        return firstHeader != null ? firstHeader.getText() : "";
+    }
 }