Merge "Moving LauncherAppWidgetHolder to dagger" into main
diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index cd8b891..d3c5c01 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -566,11 +566,14 @@
 }
 
 flag {
-  name: "google_sans_flex_font"
+  name: "gsf_res"
   namespace: "launcher"
   description: "Adds refresh for font family. Needs to be fixed to be used in resources."
   bug: "395145453"
   is_fixed_read_only: true
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
 }
 
 flag {
diff --git a/proguard.flags b/proguard.flags
index da00c00..c0a0042 100644
--- a/proguard.flags
+++ b/proguard.flags
@@ -1,6 +1,12 @@
--keep,allowshrinking,allowoptimization class com.android.launcher3.** {
-  *;
-}
+
+-keep,allowshrinking,allowoptimization class com.android.launcher3.** {*;}
+-keepclasseswithmembernames class com.android.launcher3.** {*;}
+
+-keep,allowshrinking,allowoptimization class com.android.systemui.shared.** {*;}
+-keepclasseswithmembernames class com.android.systemui.shared.** {*;}
+
+-keep,allowshrinking,allowoptimization class com.android.quickstep.** {*;}
+-keepclasseswithmembernames class com.android.quickstep.** {*;}
 
 # The support library contains references to newer platform versions.
 # Don't warn about those in case this app is linking against an older
@@ -49,14 +55,6 @@
 # Ignore warnings for hidden utility classes referenced from the shared lib
 -dontwarn com.android.internal.util.**
 
-################ Do not optimize recents lib #############
--keep class com.android.systemui.shared.** {
-  *;
-}
-
--keep class com.android.quickstep.** {
-  *;
-}
 
 -keep class com.android.internal.protolog.** {
   *;
diff --git a/quickstep/res/drawable/taskbar_divider_button_expressive_theme.xml b/quickstep/res/drawable/taskbar_divider_button_expressive_theme.xml
new file mode 100644
index 0000000..7fc61fc
--- /dev/null
+++ b/quickstep/res/drawable/taskbar_divider_button_expressive_theme.xml
@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2025 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="52dp"
+    android:height="52dp"
+    android:viewportHeight="52"
+    android:viewportWidth="52">
+    <group>
+        <path
+            android:fillColor="@color/taskbar_divider_background"
+            android:pathData="M26,14L26,38"
+            android:strokeColor="@color/taskbar_divider_background"
+            android:strokeLineCap="round"
+            android:strokeWidth="3" />
+    </group>
+</vector>
\ No newline at end of file
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index aa9e824..cb3c446 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -363,6 +363,7 @@
     <dimen name="taskbar_home_button_left_margin_kids">48dp</dimen>
     <dimen name="taskbar_icon_size_kids">32dp</dimen>
     <dimen name="taskbar_all_apps_search_button_translation_x_offset">6dp</dimen>
+    <dimen name="taskbar_all_apps_search_button_translation_x_offset_for_expressive_theme">5.5dp</dimen>
     <dimen name="taskbar_contextual_button_suw_margin">64dp</dimen>
     <dimen name="taskbar_contextual_button_suw_height">64dp</dimen>
     <dimen name="taskbar_back_button_suw_start_margin">48dp</dimen>
@@ -381,6 +382,7 @@
     <dimen name="transient_taskbar_key_shadow_distance">10dp</dimen>
     <dimen name="transient_taskbar_stashed_height">32dp</dimen>
     <dimen name="transient_taskbar_all_apps_button_translation_x_offset">8dp</dimen>
+    <dimen name="transient_taskbar_all_apps_button_translation_x_offset_for_expressive_theme">8dp</dimen>
     <dimen name="transient_taskbar_stash_spring_velocity_dp_per_s">400dp</dimen>
     <dimen name="taskbar_tooltip_vertical_padding">8dp</dimen>
     <dimen name="taskbar_tooltip_horizontal_padding">16dp</dimen>
@@ -459,7 +461,8 @@
     <!-- Container size with pointer included: bubblebar_size + bubblebar_pointer_size -->
     <dimen name="bubblebar_size_with_pointer">80dp</dimen>
     <dimen name="bubblebar_elevation">1dp</dimen>
-    <dimen name="bubblebar_drag_elevation">2dp</dimen>
+    <!-- TODO b/396539130: used wmshared value once resources are fixed -->
+    <dimen name="dragged_bubble_elevation">3dp</dimen>
     <dimen name="bubblebar_hotseat_adjustment_threshold">90dp</dimen>
     <dimen name="bubblebar_bounce_distance">20dp</dimen>
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 33cd759..7d2dc71 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -15,7 +15,6 @@
  */
 package com.android.launcher3.taskbar;
 
-import static android.view.KeyEvent.ACTION_UP;
 import static android.view.View.AccessibilityDelegate;
 import static android.view.ViewGroup.LayoutParams.MATCH_PARENT;
 import static android.view.ViewTreeObserver.InternalInsetsInfo.TOUCHABLE_INSETS_REGION;
@@ -121,6 +120,7 @@
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.StringJoiner;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.IntPredicate;
 
 /**
@@ -174,6 +174,7 @@
     private static final int NUM_ALPHA_CHANNELS = 3;
 
     private static final long AUTODIM_TIMEOUT_MS = 2250;
+    private static final long PREDICTIVE_BACK_TIMEOUT_MS = 200;
 
     private final ArrayList<StatePropertyHolder> mPropertyHolders = new ArrayList<>();
     private final ArrayList<ImageView> mAllButtons = new ArrayList<>();
@@ -894,23 +895,36 @@
     private void setBackButtonTouchListener(View buttonView,
             TaskbarNavButtonController navButtonController) {
         final RectF rect = new RectF();
+        final AtomicBoolean hasSentDownEvent = new AtomicBoolean(false);
+        final Runnable longPressTimeout = () -> {
+            navButtonController.sendBackKeyEvent(KeyEvent.ACTION_DOWN, /*cancelled*/ false);
+            hasSentDownEvent.set(true);
+        };
         buttonView.setOnTouchListener((v, event) -> {
             int motionEventAction = event.getAction();
             if (motionEventAction == MotionEvent.ACTION_DOWN) {
+                hasSentDownEvent.set(false);
+                mHandler.postDelayed(longPressTimeout, PREDICTIVE_BACK_TIMEOUT_MS);
                 rect.set(0, 0, v.getWidth(), v.getHeight());
             }
             boolean isCancelled = motionEventAction == MotionEvent.ACTION_CANCEL
                     || (!rect.contains(event.getX(), event.getY())
                     && (motionEventAction == MotionEvent.ACTION_MOVE
                     || motionEventAction == MotionEvent.ACTION_UP));
-            if (motionEventAction != MotionEvent.ACTION_DOWN
-                    && motionEventAction != MotionEvent.ACTION_UP && !isCancelled) {
-                // return early. we don't care about any other cases than DOWN, UP and CANCEL
+            if (motionEventAction != MotionEvent.ACTION_UP && !isCancelled) {
+                // return early. we don't care about any other cases than UP or CANCEL from here on
                 return false;
             }
-            int keyEventAction = motionEventAction == MotionEvent.ACTION_DOWN
-                    ? KeyEvent.ACTION_DOWN : ACTION_UP;
-            navButtonController.sendBackKeyEvent(keyEventAction, isCancelled);
+            mHandler.removeCallbacks(longPressTimeout);
+            if (!hasSentDownEvent.get()) {
+                if (isCancelled) {
+                    // if it is cancelled and ACTION_DOWN has not been sent yet, return early and
+                    // don't send anything to sysui.
+                    return false;
+                }
+                navButtonController.sendBackKeyEvent(KeyEvent.ACTION_DOWN, isCancelled);
+            }
+            navButtonController.sendBackKeyEvent(KeyEvent.ACTION_UP, isCancelled);
             if (motionEventAction == MotionEvent.ACTION_UP && !isCancelled) {
                 buttonView.performClick();
                 buttonView.performHapticFeedback(HapticFeedbackConstants.VIRTUAL_KEY);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index e41b2d2..018903e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -65,6 +65,7 @@
 import android.content.res.Resources;
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
+import android.graphics.drawable.Drawable;
 import android.hardware.display.DisplayManager;
 import android.os.IRemoteCallback;
 import android.os.Process;
@@ -98,6 +99,7 @@
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
 import com.android.launcher3.allapps.ActivityAllAppsContainerView;
 import com.android.launcher3.anim.AnimatorListeners;
 import com.android.launcher3.anim.AnimatorPlaybackController;
@@ -107,6 +109,8 @@
 import com.android.launcher3.desktop.DesktopAppLaunchTransition.AppLaunchType;
 import com.android.launcher3.folder.Folder;
 import com.android.launcher3.folder.FolderIcon;
+import com.android.launcher3.icons.BitmapRenderer;
+import com.android.launcher3.icons.FastBitmapDrawable;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.model.data.AppInfo;
@@ -141,6 +145,7 @@
 import com.android.launcher3.taskbar.customization.TaskbarSpecsEvaluator;
 import com.android.launcher3.taskbar.growth.NudgeController;
 import com.android.launcher3.taskbar.navbutton.NearestTouchFrame;
+import com.android.launcher3.taskbar.overlay.TaskbarOverlayContext;
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayController;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.shared.TestProtocol;
@@ -161,6 +166,7 @@
 import com.android.launcher3.util.TraceHelper;
 import com.android.launcher3.util.VibratorWrapper;
 import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.views.BaseDragLayer;
 import com.android.quickstep.NavHandle;
 import com.android.quickstep.RecentsModel;
 import com.android.quickstep.SystemUiProxy;
@@ -171,6 +177,7 @@
 import com.android.quickstep.views.DesktopTaskView;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
+import com.android.systemui.animation.ViewRootSync;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.rotation.RotationButtonController;
 import com.android.systemui.shared.statusbar.phone.BarTransitions;
@@ -199,7 +206,7 @@
 
     private static final String WINDOW_TITLE = "Taskbar";
 
-    private static final DesktopModeFlag ENABLE_TASKBAR_BEHIND_SHADE = new DesktopModeFlag(
+    protected static final DesktopModeFlag ENABLE_TASKBAR_BEHIND_SHADE = new DesktopModeFlag(
             Flags::enableTaskbarBehindShade, false);
 
     private final @Nullable Context mNavigationBarPanelContext;
@@ -252,6 +259,10 @@
 
     private TaskbarSpecsEvaluator mTaskbarSpecsEvaluator;
 
+    // Snapshot is used to temporarily draw taskbar behind the shade.
+    private @Nullable View mTaskbarSnapshotView;
+    private @Nullable TaskbarOverlayContext mTaskbarSnapshotOverlay;
+
     public TaskbarActivityContext(Context windowContext,
             @Nullable Context navigationBarPanelContext, DeviceProfile launcherDp,
             TaskbarNavButtonController buttonController,
@@ -309,6 +320,7 @@
         if (BubbleBarController.isBubbleBarEnabled()
                 && deviceBubbleBarEnabled
                 && bubbleBarView != null
+                && isPrimaryDisplay
         ) {
             Optional<BubbleStashedHandleViewController> bubbleHandleController = Optional.empty();
             Optional<BubbleBarSwipeController> bubbleBarSwipeController = Optional.empty();
@@ -996,6 +1008,8 @@
             mWindowManager.removeViewImmediate(mDragLayer);
             mAddedWindow = false;
         }
+        mTaskbarSnapshotView = null;
+        mTaskbarSnapshotOverlay = null;
     }
 
     public boolean isDestroyed() {
@@ -1067,6 +1081,80 @@
         if (skipAnim) {
             anim.end();
         }
+
+        updateTaskbarSnapshot(anim, isExpanded);
+    }
+
+    private void updateTaskbarSnapshot(AnimatorSet anim, boolean isExpanded) {
+        if (!ENABLE_TASKBAR_BEHIND_SHADE.isTrue()) {
+            return;
+        }
+        if (mTaskbarSnapshotView == null) {
+            mTaskbarSnapshotView = new View(this);
+        }
+        if (isExpanded) {
+            if (!mTaskbarSnapshotView.isAttachedToWindow()
+                    && mDragLayer.isAttachedToWindow()
+                    && mDragLayer.isLaidOut()
+                    && mTaskbarSnapshotView.getParent() == null) {
+                NearestTouchFrame navButtonsView = mDragLayer.findViewById(R.id.navbuttons_view);
+                int oldNavButtonsVisibility = navButtonsView.getVisibility();
+                navButtonsView.setVisibility(View.INVISIBLE);
+
+                Drawable drawable = new FastBitmapDrawable(BitmapRenderer.createHardwareBitmap(
+                        mDragLayer.getWidth(),
+                        mDragLayer.getHeight(),
+                        mDragLayer::draw));
+
+                navButtonsView.setVisibility(oldNavButtonsVisibility);
+                mTaskbarSnapshotView.setBackground(drawable);
+                mTaskbarSnapshotView.setAlpha(0f);
+
+                mTaskbarSnapshotView.addOnAttachStateChangeListener(
+                        new View.OnAttachStateChangeListener() {
+                            @Override
+                            public void onViewAttachedToWindow(@NonNull View v) {
+                                mTaskbarSnapshotView.removeOnAttachStateChangeListener(this);
+                                anim.end();
+                                mTaskbarSnapshotView.setAlpha(1f);
+                                if (!Utilities.isRunningInTestHarness()) {
+                                    ViewRootSync.synchronizeNextDraw(mDragLayer,
+                                            mTaskbarSnapshotView,
+                                            () -> {});
+                                }
+                            }
+
+                            @Override
+                            public void onViewDetachedFromWindow(@NonNull View v) {}
+                        });
+                BaseDragLayer.LayoutParams layoutParams = new BaseDragLayer.LayoutParams(
+                        mDragLayer.getWidth(), mDragLayer.getHeight());
+                layoutParams.gravity = mWindowLayoutParams.gravity;
+                layoutParams.ignoreInsets = true;
+                mTaskbarSnapshotOverlay = mControllers.taskbarOverlayController.requestWindow();
+                mTaskbarSnapshotOverlay.getDragLayer().addView(mTaskbarSnapshotView, layoutParams);
+            }
+        } else {
+            Runnable removeSnapshotView = () -> {
+                if (mTaskbarSnapshotOverlay != null) {
+                    mTaskbarSnapshotOverlay.getDragLayer().removeView(mTaskbarSnapshotView);
+                    mTaskbarSnapshotView = null;
+                    mTaskbarSnapshotOverlay = null;
+                }
+            };
+            if (mTaskbarSnapshotView.isAttachedToWindow()) {
+                mTaskbarSnapshotView.setAlpha(0f);
+                anim.end();
+                if (Utilities.isRunningInTestHarness()) {
+                    removeSnapshotView.run();
+                } else {
+                    ViewRootSync.synchronizeNextDraw(mDragLayer, mTaskbarSnapshotView,
+                            removeSnapshotView);
+                }
+            } else {
+                removeSnapshotView.run();
+            }
+        }
     }
 
     public void onRotationProposal(int rotation, boolean isValid) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index ad59c2e..eaf00b6 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -28,6 +28,7 @@
 import static com.android.launcher3.config.FeatureFlags.ENABLE_TASKBAR_NAVBAR_UNIFICATION;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TRANSIENT_TASKBAR_HIDE;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TRANSIENT_TASKBAR_SHOW;
+import static com.android.launcher3.taskbar.TaskbarActivityContext.ENABLE_TASKBAR_BEHIND_SHADE;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.launcher3.util.FlagDebugUtils.appendFlag;
 import static com.android.launcher3.util.FlagDebugUtils.formatFlagChange;
@@ -1138,7 +1139,10 @@
         long startDelay = 0;
 
         updateStateForFlag(FLAG_STASHED_IN_APP_SYSUI, hasAnyFlag(systemUiStateFlags,
-                SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE | SYSUI_STATE_DIALOG_SHOWING));
+                SYSUI_STATE_DIALOG_SHOWING | (ENABLE_TASKBAR_BEHIND_SHADE.isTrue()
+                        ? 0
+                        : SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE)
+        ));
 
         boolean stashForBubbles = hasAnyFlag(FLAG_IN_OVERVIEW)
                 && hasAnyFlag(systemUiStateFlags, SYSUI_STATE_BUBBLES_EXPANDED)
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index d43ebe2..1abef8a 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -204,7 +204,7 @@
         mExpandedBarIconsSpacing = getResources().getDimensionPixelSize(
                 R.dimen.bubblebar_expanded_icon_spacing);
         mBubbleElevation = getResources().getDimensionPixelSize(R.dimen.bubblebar_icon_elevation);
-        mDragElevation = getResources().getDimensionPixelSize(R.dimen.bubblebar_drag_elevation);
+        mDragElevation = getResources().getDimensionPixelSize(R.dimen.dragged_bubble_elevation);
         mPointerSize = getResources()
                 .getDimensionPixelSize(R.dimen.bubblebar_pointer_visible_size);
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index 7fb6480..9fb283c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -222,7 +222,7 @@
         mIconSize = res.getDimensionPixelSize(R.dimen.bubblebar_icon_size);
         mBubbleBarTaskbarMinDistance = res.getDimensionPixelSize(
                 R.dimen.bubblebar_transient_taskbar_min_distance);
-        mDragElevation = res.getDimensionPixelSize(R.dimen.bubblebar_drag_elevation);
+        mDragElevation = res.getDimensionPixelSize(R.dimen.dragged_bubble_elevation);
         mTaskbarTranslationDelta = getBubbleBarTranslationDeltaForTaskbar(activity);
         if (DeviceConfig.isSmallTablet(mActivity)) {
             mBubbleBarDropTargetSize = res.getDimensionPixelSize(R.dimen.drag_zone_bubble_fold);
diff --git a/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt
index bb2acd6..a1df21f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt
@@ -37,6 +37,7 @@
 import com.android.launcher3.views.IconButtonView
 import com.android.quickstep.DeviceConfigWrapper
 import com.android.quickstep.util.ContextualSearchStateManager
+import com.android.wm.shell.Flags
 
 /** Taskbar all apps button container for customizable taskbar. */
 class TaskbarAllAppsButtonContainer
@@ -97,14 +98,34 @@
 
     @DrawableRes
     private fun getAllAppsButton(isTransientTaskbar: Boolean): Int {
+        if (Flags.enableGsf()) {
+            return getAllAppsButtonForExpressiveTheme()
+        }
         val shouldSelectTransientIcon =
             isTransientTaskbar || (enableTaskbarPinning() && !activityContext.isThreeButtonNav)
         return if (shouldSelectTransientIcon) R.drawable.ic_transient_taskbar_all_apps_search_button
         else R.drawable.ic_taskbar_all_apps_search_button
     }
 
+    @DrawableRes
+    private fun getAllAppsButtonForExpressiveTheme(): Int {
+        return R.drawable.ic_taskbar_all_apps_search_button_expressive_theme
+    }
+
+    @DimenRes
+    fun getAllAppsButtonTranslationXOffsetForExpressiveTheme(isTransientTaskbar: Boolean): Int {
+        return if (isTransientTaskbar) {
+            R.dimen.transient_taskbar_all_apps_button_translation_x_offset_for_expressive_theme
+        } else {
+            R.dimen.taskbar_all_apps_search_button_translation_x_offset_for_expressive_theme
+        }
+    }
+
     @DimenRes
     fun getAllAppsButtonTranslationXOffset(isTransientTaskbar: Boolean): Int {
+        if (Flags.enableGsf()) {
+            return getAllAppsButtonTranslationXOffsetForExpressiveTheme(isTransientTaskbar)
+        }
         return if (isTransientTaskbar) {
             R.dimen.transient_taskbar_all_apps_button_translation_x_offset
         } else {
diff --git a/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarDividerContainer.kt b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarDividerContainer.kt
index 060ce46..08a1b48 100644
--- a/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarDividerContainer.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarDividerContainer.kt
@@ -20,6 +20,7 @@
 import android.content.Context
 import android.content.res.ColorStateList
 import android.graphics.Color.TRANSPARENT
+import android.graphics.drawable.Drawable
 import android.util.AttributeSet
 import androidx.core.view.setPadding
 import com.android.launcher3.R
@@ -28,6 +29,7 @@
 import com.android.launcher3.taskbar.TaskbarViewCallbacks
 import com.android.launcher3.views.ActivityContext
 import com.android.launcher3.views.IconButtonView
+import com.android.wm.shell.Flags
 
 /** Taskbar divider view container for customizable taskbar. */
 class TaskbarDividerContainer
@@ -46,16 +48,24 @@
         setUpIcon()
     }
 
-    @SuppressLint("UseCompatLoadingForDrawables")
     fun setUpIcon() {
         backgroundTintList = ColorStateList.valueOf(TRANSPARENT)
-        val drawable = resources.getDrawable(R.drawable.taskbar_divider_button)
+        val drawable = getTaskbarDividerIcon()
         setIconDrawable(drawable)
         if (!activityContext.isTransientTaskbar) {
             setPadding(dpToPx(activityContext.taskbarSpecsEvaluator.taskbarIconPadding.toFloat()))
         }
     }
 
+    @SuppressLint("UseCompatLoadingForDrawables")
+    fun getTaskbarDividerIcon(): Drawable {
+        return if (Flags.enableGsf()) {
+            resources.getDrawable(R.drawable.taskbar_divider_button_expressive_theme)
+        } else {
+            resources.getDrawable(R.drawable.taskbar_divider_button)
+        }
+    }
+
     @SuppressLint("ClickableViewAccessibility")
     fun setUpCallbacks(callbacks: TaskbarViewCallbacks) {
         setOnLongClickListener(callbacks.taskbarDividerLongClickListener)
diff --git a/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java b/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
index 15a27d1..e8de0d2 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
@@ -402,16 +402,16 @@
             canvas.scale(
                     mRingScale * (1f - RING_EFFECT_RATIO),
                     mRingScale * (1f - RING_EFFECT_RATIO),
-                    canvas.getWidth() / 2f,
-                    canvas.getHeight() / 2f);
+                    getWidth() / 2f,
+                    getHeight() / 2f);
         } else if (Float.compare(1, mRingScale) != 0) {
-            canvas.scale(mRingScale, mRingScale, canvas.getWidth() / 2f, canvas.getHeight() / 2f);
+            canvas.scale(mRingScale, mRingScale, getWidth() / 2f, getHeight() / 2f);
         }
         // Draw ring shadow around canvas.
         canvas.drawPath(mRingPath, mIconRingPaint);
         mIconRingPaint.setColor(mPlateColor.currentColor);
         if (Flags.enableLauncherIconShapes()) {
-            mIconRingPaint.setStrokeWidth(canvas.getWidth() * RING_EFFECT_RATIO);
+            mIconRingPaint.setStrokeWidth(getWidth() * RING_EFFECT_RATIO);
             // Using FILL_AND_STROKE as there is still some gap to fill,
             // between inner curve of ring / outer curve of icon.
             mIconRingPaint.setStyle(Paint.Style.FILL_AND_STROKE);
diff --git a/quickstep/src/com/android/quickstep/DisplayModel.kt b/quickstep/src/com/android/quickstep/DisplayModel.kt
index ac94375..3de6fd0 100644
--- a/quickstep/src/com/android/quickstep/DisplayModel.kt
+++ b/quickstep/src/com/android/quickstep/DisplayModel.kt
@@ -22,7 +22,6 @@
 import android.util.SparseArray
 import android.view.Display
 import androidx.core.util.valueIterator
-import com.android.launcher3.util.Executors
 import com.android.quickstep.DisplayModel.DisplayResource
 import java.io.PrintWriter
 
@@ -37,29 +36,24 @@
     private val displayManager = context.getSystemService(Context.DISPLAY_SERVICE) as DisplayManager
     protected val displayResourceArray = SparseArray<RESOURCE_TYPE>()
 
-    private val displayListener: DisplayManager.DisplayListener =
-        (object : DisplayManager.DisplayListener {
-            override fun onDisplayAdded(displayId: Int) {
-                if (DEBUG) Log.d(TAG, "onDisplayAdded: displayId=$displayId")
-                storeDisplayResource(displayId)
-            }
+    fun onDisplayAddSystemDecorations(displayId: Int) {
+        if (DEBUG) Log.d(TAG, "onDisplayAdded: displayId=$displayId")
+        storeDisplayResource(displayId)
+    }
 
-            override fun onDisplayRemoved(displayId: Int) {
-                if (DEBUG) Log.d(TAG, "onDisplayRemoved: displayId=$displayId")
-                deleteDisplayResource(displayId)
-            }
+    fun onDisplayRemoved(displayId: Int) {
+        if (DEBUG) Log.d(TAG, "onDisplayRemoved: displayId=$displayId")
+        deleteDisplayResource(displayId)
+    }
 
-            override fun onDisplayChanged(displayId: Int) {
-                if (DEBUG) Log.d(TAG, "onDisplayChanged: displayId=$displayId")
-            }
-        })
+    fun onDisplayRemoveSystemDecorations(displayId: Int) {
+        if (DEBUG) Log.d(TAG, "onDisplayRemoveSystemDecorations: displayId=$displayId")
+        deleteDisplayResource(displayId)
+    }
 
     protected abstract fun createDisplayResource(display: Display): RESOURCE_TYPE
 
-    protected fun registerDisplayListener() {
-        displayManager.registerDisplayListener(displayListener, Executors.MAIN_EXECUTOR.handler)
-        // In the scenario where displays were added before this display listener was
-        // registered, we should store the DisplayResources for those displays directly.
+    protected fun initializeDisplays() {
         displayManager.displays
             .filter { getDisplayResource(it.displayId) == null }
             .forEach { storeDisplayResource(it.displayId) }
@@ -70,18 +64,19 @@
             displayResource.cleanup()
         }
         displayResourceArray.clear()
-        displayManager.unregisterDisplayListener(displayListener)
     }
 
     fun getDisplayResource(displayId: Int): RESOURCE_TYPE? {
-        if (DEBUG) Log.d(TAG, "get: displayId=$displayId")
+        if (DEBUG) Log.d(TAG, Log.getStackTraceString(Throwable("get: displayId=$displayId")))
         return displayResourceArray[displayId]
     }
 
     fun deleteDisplayResource(displayId: Int) {
         if (DEBUG) Log.d(TAG, "delete: displayId=$displayId")
-        getDisplayResource(displayId)?.cleanup()
-        displayResourceArray.remove(displayId)
+        getDisplayResource(displayId)?.let {
+            it.cleanup()
+            displayResourceArray.remove(displayId)
+        }
     }
 
     fun storeDisplayResource(displayId: Int) {
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
index 152630a..6a9c3dd 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
@@ -91,27 +91,33 @@
      */
     private var keyboardTaskFocusIndex = -1
 
-    // TODO (b/397942185): get per-display interface
-    private val containerInterface: BaseContainerInterface<*, *>
-        get() = overviewComponentObserver.getContainerInterface(DEFAULT_DISPLAY)
+    private fun getContainerInterface(displayId: Int) =
+        overviewComponentObserver.getContainerInterface(displayId)
 
-    // TODO (b/397942185): get per-display RecentsView
-    private val visibleRecentsView: RecentsView<*, *>?
-        get() = containerInterface.getVisibleRecentsView<RecentsView<*, *>>()
+    private fun getVisibleRecentsView(displayId: Int) =
+        getContainerInterface(displayId).getVisibleRecentsView<RecentsView<*, *>>()
 
     /**
      * Adds a command to be executed next, after all pending tasks are completed. Max commands that
      * can be queued is [.MAX_QUEUE_SIZE]. Requests after reaching that limit will be silently
      * dropped.
+     *
+     * @param type The type of the command
+     * @param onDisplays The display to run the command on
      */
     @BinderThread
-    fun addCommand(type: CommandType): CommandInfo? {
+    @JvmOverloads
+    fun addCommand(
+        type: CommandType,
+        displayId: Int = DEFAULT_DISPLAY,
+        isLastOfBatch: Boolean = true,
+    ): CommandInfo? {
         if (commandQueue.size >= MAX_QUEUE_SIZE) {
             Log.d(TAG, "command not added: $type - queue is full ($commandQueue).")
             return null
         }
 
-        val command = CommandInfo(type)
+        val command = CommandInfo(type, displayId = displayId, isLastOfBatch = isLastOfBatch)
         commandQueue.add(command)
         Log.d(TAG, "command added: $command")
 
@@ -129,6 +135,35 @@
         return command
     }
 
+    @BinderThread
+    fun addCommandsForDisplays(type: CommandType, displayIds: IntArray): CommandInfo? {
+        if (displayIds.isEmpty()) return null
+        var lastCommand: CommandInfo? = null
+        displayIds.forEachIndexed({ i, displayId ->
+            lastCommand = addCommand(type, displayId, i == displayIds.size - 1)
+        })
+        return lastCommand
+    }
+
+    @BinderThread
+    fun addCommandsForAllDisplays(type: CommandType) =
+        addCommandsForDisplays(
+            type,
+            recentsDisplayModel.activeDisplayResources
+                .map { resource -> resource.displayId }
+                .toIntArray(),
+        )
+
+    @BinderThread
+    fun addCommandsForDisplaysExcept(type: CommandType, excludedDisplayId: Int) =
+        addCommandsForDisplays(
+            type,
+            recentsDisplayModel.activeDisplayResources
+                .map { resource -> resource.displayId }
+                .filter { displayId -> displayId != excludedDisplayId }
+                .toIntArray(),
+        )
+
     fun canStartHomeSafely(): Boolean = commandQueue.isEmpty() || commandQueue.first().type == HOME
 
     /** Clear pending or completed commands from the queue */
@@ -143,7 +178,7 @@
      * completion (returns false).
      */
     @UiThread
-    private fun processNextCommand() =
+    private fun processNextCommand(): Unit =
         traceSection("OverviewCommandHelper.processNextCommand") {
             val command: CommandInfo? = commandQueue.firstOrNull()
             if (command == null) {
@@ -182,7 +217,7 @@
      */
     @VisibleForTesting
     fun executeCommand(command: CommandInfo, onCallbackResult: () -> Unit): Boolean {
-        val recentsView = visibleRecentsView
+        val recentsView = getVisibleRecentsView(command.displayId)
         Log.d(TAG, "executeCommand: $command - visibleRecentsView: $recentsView")
         return if (recentsView != null) {
             executeWhenRecentsIsVisible(command, recentsView, onCallbackResult)
@@ -230,6 +265,7 @@
                     launchTask(recentsView, taskView, command, onCallbackResult)
                 }
             }
+
             TOGGLE -> {
                 launchTask(
                     recentsView,
@@ -238,6 +274,7 @@
                     onCallbackResult,
                 )
             }
+
             HOME -> {
                 recentsView.startHome()
                 true
@@ -294,6 +331,7 @@
         command: CommandInfo,
         onCallbackResult: () -> Unit,
     ): Boolean {
+        val containerInterface = getContainerInterface(command.displayId)
         val recentsViewContainer = containerInterface.getCreatedContainer()
         val recentsView: RecentsView<*, *>? = recentsViewContainer?.getOverviewPanel()
         val deviceProfile = recentsViewContainer?.getDeviceProfile()
@@ -335,6 +373,7 @@
 
                 if (keyboardTaskFocusIndex == -1) return true
             }
+
             KEYBOARD_INPUT ->
                 if (uiController != null && deviceProfile?.isTablet == true) {
                     if (
@@ -348,6 +387,7 @@
                 } else {
                     keyboardTaskFocusIndex = 0
                 }
+
             HOME -> {
                 ActiveGestureProtoLogProxy.logExecuteHomeCommand()
                 // Although IActivityTaskManager$Stub$Proxy.startActivity is a slow binder call,
@@ -357,12 +397,14 @@
                 touchInteractionService.startActivity(overviewComponentObserver.homeIntent)
                 return true
             }
+
             SHOW ->
                 // When Recents is not currently visible, the command's type is SHOW
                 // when overview is triggered via the keyboard overview button or Action+Tab
                 // keys (Not Alt+Tab which is KQS). The overview button on-screen in 3-button
                 // nav is TYPE_TOGGLE.
                 keyboardTaskFocusIndex = 0
+
             TOGGLE -> {}
         }
 
@@ -378,7 +420,7 @@
                     Log.d(TAG, "switching to Overview state - onAnimationStart: $command")
                     super.onAnimationStart(animation)
                     updateRecentsViewFocus(command)
-                    logShowOverviewFrom(command.type)
+                    logShowOverviewFrom(command)
                 }
 
                 override fun onAnimationEnd(animation: Animator) {
@@ -402,7 +444,7 @@
 
         val gestureState =
             touchInteractionService.createGestureState(
-                focusedDisplayId,
+                command.displayId,
                 GestureState.DEFAULT_STATE,
                 GestureState.TrackpadGestureType.NONE,
             )
@@ -432,7 +474,7 @@
                     }
 
                     updateRecentsViewFocus(command)
-                    logShowOverviewFrom(command.type)
+                    logShowOverviewFrom(command)
                     containerInterface.runOnInitBackgroundStateUI {
                         Log.d(TAG, "recents animation started - onInitBackgroundStateUI: $command")
                         interactionHandler.onGestureEnded(
@@ -456,12 +498,13 @@
                 }
             }
 
-        val displayId = gestureState.displayId
         val taskAnimationManager =
-            recentsDisplayModel.getTaskAnimationManager(displayId)
+            recentsDisplayModel.getTaskAnimationManager(command.displayId)
                 ?: run {
-                    Log.e(TAG, "No TaskAnimationManager found for display $displayId")
-                    ActiveGestureProtoLogProxy.logOnTaskAnimationManagerNotAvailable(displayId)
+                    Log.e(TAG, "No TaskAnimationManager found for display ${command.displayId}")
+                    ActiveGestureProtoLogProxy.logOnTaskAnimationManagerNotAvailable(
+                        command.displayId
+                    )
                     return false
                 }
         if (taskAnimationManager.isRecentsAnimationRunning) {
@@ -526,8 +569,13 @@
     }
 
     private fun updateRecentsViewFocus(command: CommandInfo) {
-        val recentsView: RecentsView<*, *> = visibleRecentsView ?: return
-        if (command.type != KEYBOARD_INPUT && command.type != HIDE && command.type != SHOW) {
+        val recentsView: RecentsView<*, *> = getVisibleRecentsView(command.displayId) ?: return
+        if (
+            command.type != KEYBOARD_INPUT &&
+                command.type != HIDE &&
+                command.type != SHOW &&
+                command.type != TOGGLE
+        ) {
             return
         }
 
@@ -547,7 +595,7 @@
     }
 
     private fun onRecentsViewFocusUpdated(command: CommandInfo) {
-        val recentsView: RecentsView<*, *> = visibleRecentsView ?: return
+        val recentsView: RecentsView<*, *> = getVisibleRecentsView(command.displayId) ?: return
         if (command.type != HIDE || keyboardTaskFocusIndex == PagedView.INVALID_PAGE) {
             return
         }
@@ -565,10 +613,11 @@
         return true
     }
 
-    private fun logShowOverviewFrom(commandType: CommandType) {
+    private fun logShowOverviewFrom(command: CommandInfo) {
+        val containerInterface = getContainerInterface(command.displayId)
         val container = containerInterface.getCreatedContainer() ?: return
         val event =
-            when (commandType) {
+            when (command.type) {
                 SHOW -> LAUNCHER_OVERVIEW_SHOW_OVERVIEW_FROM_KEYBOARD_SHORTCUT
                 HIDE -> LAUNCHER_OVERVIEW_SHOW_OVERVIEW_FROM_KEYBOARD_QUICK_SWITCH
                 TOGGLE -> LAUNCHER_OVERVIEW_SHOW_OVERVIEW_FROM_3_BUTTON
@@ -601,6 +650,8 @@
         var status: CommandStatus = CommandStatus.IDLE,
         val createTime: Long = SystemClock.elapsedRealtime(),
         private var animationCallbacks: RecentsAnimationCallbacks? = null,
+        val displayId: Int = DEFAULT_DISPLAY,
+        val isLastOfBatch: Boolean = true,
     ) {
         fun setAnimationCallbacks(recentsAnimationCallbacks: RecentsAnimationCallbacks) {
             this.animationCallbacks = recentsAnimationCallbacks
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 6912ba7..7878e68 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -316,7 +316,10 @@
         @Override
         public void onDisplayAddSystemDecorations(int displayId) {
             executeForTaskbarManager(taskbarManager ->
-                            taskbarManager.onDisplayAddSystemDecorations(displayId));
+                    taskbarManager.onDisplayAddSystemDecorations(displayId));
+
+            executeForRecentsDisplayModel(displayModel ->
+                    displayModel.onDisplayAddSystemDecorations(displayId));
         }
 
         @BinderThread
@@ -327,6 +330,8 @@
             executeForTouchInteractionService(tis -> {
                 tis.mDeviceState.clearSysUIStateFlagsForDisplay(displayId);
             });
+            executeForRecentsDisplayModel(displayModel ->
+                    displayModel.onDisplayRemoved(displayId));
         }
 
         @BinderThread
@@ -334,6 +339,8 @@
         public void onDisplayRemoveSystemDecorations(int displayId) {
             executeForTaskbarManager(taskbarManager ->
                     taskbarManager.onDisplayRemoveSystemDecorations(displayId));
+            executeForRecentsDisplayModel(displayModel ->
+                    displayModel.onDisplayRemoveSystemDecorations(displayId));
         }
 
         @BinderThread
@@ -445,6 +452,15 @@
             }));
         }
 
+        private void executeForRecentsDisplayModel(
+                @NonNull Consumer<RecentsDisplayModel> recentsDisplayModelConsumer) {
+            MAIN_EXECUTOR.execute(() -> executeForTouchInteractionService(tis -> {
+                RecentsDisplayModel recentsDisplayModel = tis.mRecentsDisplayModel;
+                if (recentsDisplayModel == null) return;
+                recentsDisplayModelConsumer.accept(recentsDisplayModel);
+            }));
+        }
+
         /**
          * Returns the {@link TaskbarManager}.
          * <p>
@@ -1251,7 +1267,7 @@
 
         private InputMonitorDisplayModel(Context context) {
             super(context);
-            registerDisplayListener();
+            initializeDisplays();
         }
 
         @NonNull
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt
index 0f611eb..5b88686 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt
@@ -17,6 +17,7 @@
 package com.android.quickstep.fallback.window
 
 import android.content.Context
+import android.util.Log
 import android.view.Display
 import android.view.Display.DEFAULT_DISPLAY
 import androidx.core.util.valueIterator
@@ -57,7 +58,7 @@
 
     init {
         if (enableOverviewInWindow) {
-            registerDisplayListener()
+            initializeDisplays()
         } else {
             // Always create resource for default display
             storeDisplayResource(DEFAULT_DISPLAY)
@@ -74,10 +75,12 @@
     }
 
     fun getRecentsWindowManager(displayId: Int): RecentsWindowManager? {
+        if (DEBUG) Log.d(TAG, "getRecentsWindowManager for display $displayId")
         return getDisplayResource(displayId)?.recentsWindowManager
     }
 
     fun getFallbackWindowInterface(displayId: Int): FallbackWindowInterface? {
+        if (DEBUG) Log.d(TAG, "getFallbackWindowInterface for display $displayId")
         return getDisplayResource(displayId)?.fallbackWindowInterface
     }
 
diff --git a/quickstep/src/com/android/quickstep/util/DesksUtils.kt b/quickstep/src/com/android/quickstep/util/DesksUtils.kt
index a10f771..521ba27 100644
--- a/quickstep/src/com/android/quickstep/util/DesksUtils.kt
+++ b/quickstep/src/com/android/quickstep/util/DesksUtils.kt
@@ -19,7 +19,6 @@
 import android.app.TaskInfo
 import android.content.ComponentName
 import android.content.res.Resources
-import android.util.Log
 import android.window.DesktopExperienceFlags
 import com.android.systemui.shared.recents.model.Task
 
@@ -40,7 +39,9 @@
 
         @JvmStatic
         fun isDesktopWallpaperTask(taskInfo: TaskInfo): Boolean {
-            Log.d("b/403118101", "isDesktopWallpaperTask: $taskInfo")
+            // TODO: b/403118101 - In some launcher tests, there is a task with baseIntent set to
+            // null. Remove this check after finding out how that task is created.
+            if (taskInfo.baseIntent == null) return false
             return taskInfo.baseIntent.component?.let(::isDesktopWallpaperComponent) == true
         }
 
diff --git a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
index d2f9652..54ac1cf 100644
--- a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
+++ b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
@@ -137,7 +137,8 @@
         mDestinationBoundsTransformed.set(destinationBoundsTransformed);
         mSurfaceTransactionHelper = new PipSurfaceTransactionHelper(cornerRadius, shadowRadius);
 
-        final float aspectRatio = destinationBounds.width() / (float) destinationBounds.height();
+        final Rational aspectRatio = new Rational(
+                destinationBounds.width(), destinationBounds.height());
         String reasonForCreateOverlay = null; // For debugging purpose.
 
         // Slightly larger app bounds to allow for off by 1 pixel source-rect-hint errors.
@@ -158,16 +159,22 @@
             // not a valid rectangle to use for cropping app surface
             reasonForCreateOverlay = "Source rect hint exceeds display bounds " + sourceRectHint;
             sourceRectHint.setEmpty();
-        } else if (!PictureInPictureParams.isSameAspectRatio(sourceRectHint,
-                new Rational(destinationBounds.width(), destinationBounds.height()))) {
-            // The source rect hint does not aspect ratio
-            reasonForCreateOverlay = "Source rect hint does not match aspect ratio "
-                    + sourceRectHint + " aspect ratio " + aspectRatio;
-            sourceRectHint.setEmpty();
+        } else {
+            final Rational srcAspectRatio = new Rational(
+                    sourceRectHint.width(), sourceRectHint.height());
+            if (!PictureInPictureParams.isSameAspectRatio(destinationBounds, srcAspectRatio)) {
+                // The aspect ratio of destination bounds does not match source rect hint.
+                // We use the aspect ratio of source rect hint to check against destination bounds
+                // here to avoid upscaling error.
+                reasonForCreateOverlay = "Source rect hint:" + sourceRectHint
+                        + " does not match destination bounds:" + destinationBounds;
+                sourceRectHint.setEmpty();
+            }
         }
 
         if (sourceRectHint.isEmpty()) {
-            mSourceRectHint.set(getEnterPipWithOverlaySrcRectHint(appBounds, aspectRatio));
+            mSourceRectHint.set(
+                    getEnterPipWithOverlaySrcRectHint(appBounds, aspectRatio.floatValue()));
             mPipContentOverlay = new PipContentOverlay.PipAppIconOverlay(view.getContext(),
                     mAppBounds, mDestinationBounds,
                     new IconProvider(context).getIcon(mActivityInfo), appIconSizePx);
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
index 381ac68..11e0ee8 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
@@ -17,6 +17,7 @@
 package com.android.quickstep
 
 import android.platform.test.flag.junit.SetFlagsRule
+import android.view.Display.DEFAULT_DISPLAY
 import androidx.test.filters.SmallTest
 import com.android.launcher3.Flags
 import com.android.launcher3.util.LauncherMultivalentJUnit
@@ -25,6 +26,7 @@
 import com.android.quickstep.OverviewCommandHelper.CommandInfo
 import com.android.quickstep.OverviewCommandHelper.CommandInfo.CommandStatus
 import com.android.quickstep.OverviewCommandHelper.CommandType
+import com.android.quickstep.fallback.window.RecentsDisplayModel
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.delay
@@ -41,9 +43,9 @@
 import org.junit.runner.RunWith
 import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.spy
-import org.mockito.Mockito.`when`
 import org.mockito.kotlin.any
 import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @SmallTest
 @RunWith(LauncherMultivalentJUnit::class)
@@ -57,18 +59,38 @@
 
     private var pendingCallbacksWithDelays = mutableListOf<Long>()
 
+    private val recentsDisplayModel: RecentsDisplayModel = mock()
+    private val defaultDisplayResource: RecentsDisplayModel.RecentsDisplayResource = mock()
+    private val secondaryDisplayResource: RecentsDisplayModel.RecentsDisplayResource = mock()
+    private val executeCommandDisplayIds = mutableListOf<Int>()
+
+    private fun setupDefaultDisplay() {
+        whenever(defaultDisplayResource.displayId).thenReturn(DEFAULT_DISPLAY)
+        whenever(recentsDisplayModel.activeDisplayResources)
+            .thenReturn(listOf(defaultDisplayResource))
+    }
+
+    private fun setupMultipleDisplays() {
+        whenever(defaultDisplayResource.displayId).thenReturn(DEFAULT_DISPLAY)
+        whenever(secondaryDisplayResource.displayId).thenReturn(1)
+        whenever(recentsDisplayModel.activeDisplayResources)
+            .thenReturn(listOf(defaultDisplayResource, secondaryDisplayResource))
+    }
+
     @Suppress("UNCHECKED_CAST")
     @Before
     fun setup() {
         setFlagsRule.setFlags(true, Flags.FLAG_ENABLE_OVERVIEW_COMMAND_HELPER_TIMEOUT)
 
+        setupDefaultDisplay()
+
         sut =
             spy(
                 OverviewCommandHelper(
                     touchInteractionService = mock(),
                     overviewComponentObserver = mock(),
                     dispatcherProvider = TestDispatcherProvider(dispatcher),
-                    recentsDisplayModel = mock(),
+                    recentsDisplayModel = recentsDisplayModel,
                     focusState = mock(),
                     taskbarManager = mock(),
                 )
@@ -86,6 +108,8 @@
                         }
                     }
                 }
+                val commandInfo = invocation.arguments[0] as CommandInfo
+                executeCommandDisplayIds.add(commandInfo.displayId)
                 delayInMillis == null // if no callback to execute, returns success
             }
             .`when`(sut)
@@ -175,7 +199,61 @@
             assertThat(commandInfo2.status).isEqualTo(CommandStatus.COMPLETED)
         }
 
+    @Test
+    fun whenAllDisplaysCommandIsAdded_singleCommandProcessedForDefaultDisplay() =
+        testScope.runTest {
+            executeCommandDisplayIds.clear()
+            // Add command to queue
+            val commandInfo: CommandInfo = sut.addCommandsForAllDisplays(CommandType.HOME)!!
+            assertThat(commandInfo.status).isEqualTo(CommandStatus.IDLE)
+            runCurrent()
+            assertThat(commandInfo.status).isEqualTo(CommandStatus.COMPLETED)
+            assertThat(executeCommandDisplayIds).containsExactly(DEFAULT_DISPLAY)
+        }
+
+    @Test
+    fun whenAllDisplaysCommandIsAdded_multipleCommandsProcessedForMultipleDisplays() =
+        testScope.runTest {
+            setupMultipleDisplays()
+            executeCommandDisplayIds.clear()
+            // Add command to queue
+            val commandInfo: CommandInfo = sut.addCommandsForAllDisplays(CommandType.HOME)!!
+            assertThat(commandInfo.status).isEqualTo(CommandStatus.IDLE)
+            runCurrent()
+            assertThat(commandInfo.status).isEqualTo(CommandStatus.COMPLETED)
+            assertThat(executeCommandDisplayIds)
+                .containsExactly(DEFAULT_DISPLAY, EXTERNAL_DISPLAY_ID)
+        }
+
+    @Test
+    fun whenAllExceptDisplayCommandIsAdded_otherDisplayProcessed() =
+        testScope.runTest {
+            setupMultipleDisplays()
+            executeCommandDisplayIds.clear()
+            // Add command to queue
+            val commandInfo: CommandInfo =
+                sut.addCommandsForDisplaysExcept(CommandType.HOME, DEFAULT_DISPLAY)!!
+            assertThat(commandInfo.status).isEqualTo(CommandStatus.IDLE)
+            runCurrent()
+            assertThat(commandInfo.status).isEqualTo(CommandStatus.COMPLETED)
+            assertThat(executeCommandDisplayIds).containsExactly(EXTERNAL_DISPLAY_ID)
+        }
+
+    @Test
+    fun whenSingleDisplayCommandIsAdded_thatDisplayIsProcessed() =
+        testScope.runTest {
+            executeCommandDisplayIds.clear()
+            val displayId = 5
+            // Add command to queue
+            val commandInfo: CommandInfo = sut.addCommand(CommandType.HOME, displayId)!!
+            assertThat(commandInfo.status).isEqualTo(CommandStatus.IDLE)
+            runCurrent()
+            assertThat(commandInfo.status).isEqualTo(CommandStatus.COMPLETED)
+            assertThat(executeCommandDisplayIds).containsExactly(displayId)
+        }
+
     private companion object {
         const val QUEUE_TIMEOUT = 5001L
+        const val EXTERNAL_DISPLAY_ID = 1
     }
 }
diff --git a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
index 7b73be7..746f8bb 100644
--- a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
@@ -22,6 +22,7 @@
 import android.content.pm.PackageManager
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import android.view.Display.DEFAULT_DISPLAY
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
@@ -52,6 +53,7 @@
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.mockito.Mockito.`when`
 import org.mockito.kotlin.any
@@ -64,8 +66,11 @@
 import org.mockito.quality.Strictness
 
 /** Test for [DesktopSystemShortcut] */
+// TODO(b/403558856): Improve test coverage for DesktopModeCompatPolicy integration.
 class DesktopSystemShortcutTest {
 
+    @get:Rule val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
+
     private val launcher: RecentsViewContainer = mock()
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogManager.StatsLogger = mock()
diff --git a/res/drawable/ic_taskbar_all_apps_search_button_expressive_theme.xml b/res/drawable/ic_taskbar_all_apps_search_button_expressive_theme.xml
new file mode 100644
index 0000000..ed4a821
--- /dev/null
+++ b/res/drawable/ic_taskbar_all_apps_search_button_expressive_theme.xml
@@ -0,0 +1,25 @@
+<!--
+  ~ Copyright (C) 2025 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="52dp"
+    android:height="52dp"
+    android:viewportWidth="52"
+    android:viewportHeight="52">
+    <path
+        android:pathData="M32.862,38.75C31.263,38.75 29.888,38.175 28.737,37.025C27.587,35.875 27.013,34.487 27.013,32.862C27.013,31.237 27.587,29.85 28.737,28.7C29.888,27.55 31.263,26.975 32.862,26.975C34.487,26.975 35.875,27.55 37.025,28.7C38.175,29.85 38.75,31.237 38.75,32.862C38.75,33.438 38.662,33.987 38.487,34.513C38.338,35.037 38.125,35.525 37.85,35.975L41.45,39.612C41.6,39.737 41.7,39.875 41.75,40.025C41.825,40.175 41.862,40.338 41.862,40.513C41.862,40.688 41.825,40.85 41.75,41C41.7,41.15 41.6,41.3 41.45,41.45C41.325,41.575 41.188,41.675 41.037,41.75C40.888,41.8 40.725,41.825 40.55,41.825C40.375,41.825 40.213,41.8 40.063,41.75C39.912,41.675 39.763,41.575 39.612,41.45L35.975,37.85C35.525,38.15 35.037,38.375 34.513,38.525C33.987,38.675 33.438,38.75 32.862,38.75ZM16.1,38.75C14.5,38.75 13.125,38.175 11.975,37.025C10.825,35.875 10.25,34.5 10.25,32.9C10.25,31.25 10.825,29.85 11.975,28.7C13.125,27.55 14.5,26.975 16.1,26.975C17.75,26.975 19.15,27.55 20.3,28.7C21.45,29.85 22.025,31.25 22.025,32.9C22.025,34.5 21.45,35.875 20.3,37.025C19.15,38.175 17.75,38.75 16.1,38.75ZM32.862,36.125C33.763,36.125 34.525,35.813 35.15,35.188C35.8,34.537 36.125,33.775 36.125,32.9C36.125,31.975 35.8,31.2 35.15,30.575C34.525,29.925 33.763,29.6 32.862,29.6C31.962,29.6 31.188,29.925 30.538,30.575C29.913,31.2 29.6,31.975 29.6,32.9C29.6,33.775 29.913,34.537 30.538,35.188C31.188,35.813 31.962,36.125 32.862,36.125ZM16.1,22.025C14.5,22.025 13.125,21.45 11.975,20.3C10.825,19.15 10.25,17.763 10.25,16.138C10.25,14.512 10.825,13.125 11.975,11.975C13.125,10.825 14.5,10.25 16.1,10.25C17.75,10.25 19.15,10.825 20.3,11.975C21.45,13.125 22.025,14.512 22.025,16.138C22.025,17.763 21.45,19.15 20.3,20.3C19.15,21.45 17.75,22.025 16.1,22.025ZM32.862,22.025C31.237,22.025 29.85,21.45 28.7,20.3C27.55,19.15 26.975,17.763 26.975,16.138C26.975,14.512 27.55,13.125 28.7,11.975C29.85,10.825 31.237,10.25 32.862,10.25C34.487,10.25 35.875,10.825 37.025,11.975C38.175,13.125 38.75,14.512 38.75,16.138C38.75,17.763 38.175,19.15 37.025,20.3C35.875,21.45 34.487,22.025 32.862,22.025Z"
+        android:fillColor="#FF000000"/>
+</vector>
\ No newline at end of file
diff --git a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
index f223eaa..bf02e03 100644
--- a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
+++ b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
@@ -435,6 +435,7 @@
                 }
                 if (currentItem.itemInfo != null && Objects.equals(
                         currentItem.itemInfo.getTargetPackage(), PRIVATE_SPACE_PACKAGE)) {
+                    currentItem.itemInfo.bitmap = mPrivateProviderManager.preparePSBitmapInfo();
                     currentItem.itemInfo.bitmap.creationFlags |= FLAG_NO_BADGE;
                     currentItem.itemInfo.contentDescription =
                             mPrivateProviderManager.getPsAppContentDesc();
diff --git a/src/com/android/launcher3/allapps/PrivateProfileManager.java b/src/com/android/launcher3/allapps/PrivateProfileManager.java
index 1bc1b17..0e6a5b8 100644
--- a/src/com/android/launcher3/allapps/PrivateProfileManager.java
+++ b/src/com/android/launcher3/allapps/PrivateProfileManager.java
@@ -190,15 +190,11 @@
     /** Adds Private Space install app button to the layout. */
     public void addPrivateSpaceInstallAppButton(List<BaseAllAppsAdapter.AdapterItem> adapterItems) {
         Context context = mAllApps.getContext();
-        // Prepare bitmapInfo
-        Intent.ShortcutIconResource shortcut = Intent.ShortcutIconResource.fromContext(
-                context, com.android.launcher3.R.drawable.private_space_install_app_icon);
-        BitmapInfo bitmapInfo = LauncherIcons.obtain(context).createIconBitmap(shortcut);
 
         PrivateSpaceInstallAppButtonInfo itemInfo = new PrivateSpaceInstallAppButtonInfo();
         itemInfo.title = context.getResources().getString(R.string.ps_add_button_label);
         itemInfo.intent = mAppInstallerIntent;
-        itemInfo.bitmap = bitmapInfo;
+        itemInfo.bitmap = preparePSBitmapInfo();
         itemInfo.contentDescription = context.getResources().getString(
                 com.android.launcher3.R.string.ps_add_button_content_description);
         itemInfo.runtimeStatusFlags |= FLAG_NOT_PINNABLE;
@@ -218,6 +214,13 @@
                     .get(mAllApps.getContext()).getValue(PRIVATE_SPACE_HIDE_WHEN_LOCKED_URI, 0);
     }
 
+    BitmapInfo preparePSBitmapInfo() {
+        Context context = mAllApps.getContext();
+        Intent.ShortcutIconResource shortcut = Intent.ShortcutIconResource.fromContext(
+                context, com.android.launcher3.R.drawable.private_space_install_app_icon);
+        return LauncherIcons.obtain(context).createIconBitmap(shortcut);
+    }
+
     /**
      * Resets the current state of Private Profile, w.r.t. to Launcher. The decorator should only
      * be applied upon expand before animating. When collapsing, reset() will remove the decorator
diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProxy.java b/src/com/android/launcher3/graphics/GridCustomizationsProxy.java
index 48519ce..b40e099 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProxy.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProxy.java
@@ -143,7 +143,7 @@
     private final InvariantDeviceProfile mIdp;
 
     @Inject
-    GridCustomizationsProxy(
+    protected GridCustomizationsProxy(
             @ApplicationContext Context context,
             ThemeManager themeManager,
             LauncherPrefs prefs,