Merge "Restructure camera platform flags" into main
diff --git a/core/java/android/window/TaskConstants.java b/core/java/android/window/TaskConstants.java
index 69d79b4..44bb33db 100644
--- a/core/java/android/window/TaskConstants.java
+++ b/core/java/android/window/TaskConstants.java
@@ -81,6 +81,12 @@
     public static final int TASK_CHILD_LAYER_RESIZE_VEIL = 6 * TASK_CHILD_LAYER_REGION_SIZE;
 
     /**
+     * Floating menus belonging to a task (e.g. maximize menu).
+     * @hide
+     */
+    public  static final int TASK_CHILD_LAYER_FLOATING_MENU = 7 * TASK_CHILD_LAYER_REGION_SIZE;
+
+    /**
      * Z-orders of task child layers other than activities, task fragments and layers interleaved
      * with them, e.g. IME windows. [-10000, 10000) is reserved for these layers.
      * @hide
diff --git a/libs/WindowManager/Shell/res/color/desktop_mode_maximize_menu_button_color_selector.xml b/libs/WindowManager/Shell/res/color/desktop_mode_maximize_menu_button_color_selector.xml
new file mode 100644
index 0000000..65f5239
--- /dev/null
+++ b/libs/WindowManager/Shell/res/color/desktop_mode_maximize_menu_button_color_selector.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2023 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License
+-->
+
+<selector xmlns:android="http://schemas.android.com/apk/res/android">
+    <item android:state_pressed="true"
+        android:color="@color/desktop_mode_maximize_menu_button_on_hover"/>
+    <item android:state_hovered="true"
+        android:color="@color/desktop_mode_maximize_menu_button_on_hover"/>
+    <item android:state_focused="true"
+        android:color="@color/desktop_mode_maximize_menu_button_on_hover"/>
+    <item android:state_selected="true"
+        android:color="@color/desktop_mode_maximize_menu_button_on_hover"/>
+    <item android:color="@color/desktop_mode_maximize_menu_button"/>
+</selector>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/res/color/desktop_mode_maximize_menu_button_outline_color_selector.xml b/libs/WindowManager/Shell/res/color/desktop_mode_maximize_menu_button_outline_color_selector.xml
new file mode 100644
index 0000000..86679af
--- /dev/null
+++ b/libs/WindowManager/Shell/res/color/desktop_mode_maximize_menu_button_outline_color_selector.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2023 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License
+-->
+
+<selector xmlns:android="http://schemas.android.com/apk/res/android">
+    <item android:state_pressed="true"
+        android:color="@color/desktop_mode_maximize_menu_button_outline_on_hover"/>
+    <item android:state_hovered="true"
+        android:color="@color/desktop_mode_maximize_menu_button_outline_on_hover"/>
+    <item android:state_focused="true"
+        android:color="@color/desktop_mode_maximize_menu_button_outline_on_hover"/>
+    <item android:state_selected="true"
+        android:color="@color/desktop_mode_maximize_menu_button_outline_on_hover"/>
+    <item android:color="@color/desktop_mode_maximize_menu_button_outline"/>
+</selector>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/res/drawable/desktop_mode_decor_menu_background.xml b/libs/WindowManager/Shell/res/drawable/desktop_mode_decor_handle_menu_background.xml
similarity index 100%
rename from libs/WindowManager/Shell/res/drawable/desktop_mode_decor_menu_background.xml
rename to libs/WindowManager/Shell/res/drawable/desktop_mode_decor_handle_menu_background.xml
diff --git a/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_background.xml b/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_background.xml
new file mode 100644
index 0000000..5d9fe67
--- /dev/null
+++ b/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_background.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2023 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<shape android:shape="rectangle"
+    xmlns:android="http://schemas.android.com/apk/res/android">
+    <solid android:color="@android:color/white" />
+    <corners android:radius="@dimen/desktop_mode_maximize_menu_corner_radius" />
+</shape>
diff --git a/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_maximize_button_background.xml b/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_maximize_button_background.xml
new file mode 100644
index 0000000..bfb0dd7
--- /dev/null
+++ b/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_maximize_button_background.xml
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2023 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License
+-->
+
+<shape xmlns:android="http://schemas.android.com/apk/res/android"
+    android:shape="rectangle">
+    <solid android:color="@color/desktop_mode_maximize_menu_button_color_selector"/>
+    <corners
+        android:radius="@dimen/desktop_mode_maximize_menu_buttons_large_corner_radius"/>
+    <stroke android:width="1dp" android:color="@color/desktop_mode_maximize_menu_button_outline_color_selector"/>
+</shape>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_snap_left_button_background.xml b/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_snap_left_button_background.xml
new file mode 100644
index 0000000..6630fca
--- /dev/null
+++ b/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_snap_left_button_background.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+
+<!--
+  ~ Copyright (C) 2023 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License
+-->
+
+<shape xmlns:android="http://schemas.android.com/apk/res/android"
+    android:shape="rectangle">
+    <solid android:color="@color/desktop_mode_maximize_menu_button_color_selector"/>
+    <corners
+        android:topLeftRadius="@dimen/desktop_mode_maximize_menu_buttons_large_corner_radius"
+        android:topRightRadius="@dimen/desktop_mode_maximize_menu_buttons_small_corner_radius"
+        android:bottomLeftRadius="@dimen/desktop_mode_maximize_menu_buttons_large_corner_radius"
+        android:bottomRightRadius="@dimen/desktop_mode_maximize_menu_buttons_small_corner_radius"/>
+    <stroke android:width="1dp" android:color="@color/desktop_mode_maximize_menu_button_outline_color_selector"/>
+</shape>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_snap_right_button_background.xml b/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_snap_right_button_background.xml
new file mode 100644
index 0000000..7bd6e99
--- /dev/null
+++ b/libs/WindowManager/Shell/res/drawable/desktop_mode_maximize_menu_snap_right_button_background.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+
+<!--
+  ~ Copyright (C) 2023 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License
+-->
+
+<shape xmlns:android="http://schemas.android.com/apk/res/android"
+    android:shape="rectangle">
+    <solid android:color="@color/desktop_mode_maximize_menu_button_color_selector"/>
+    <corners
+        android:topLeftRadius="@dimen/desktop_mode_maximize_menu_buttons_small_corner_radius"
+        android:topRightRadius="@dimen/desktop_mode_maximize_menu_buttons_large_corner_radius"
+        android:bottomLeftRadius="@dimen/desktop_mode_maximize_menu_buttons_small_corner_radius"
+        android:bottomRightRadius="@dimen/desktop_mode_maximize_menu_buttons_large_corner_radius"/>
+    <stroke android:width="1dp" android:color="@color/desktop_mode_maximize_menu_button_outline_color_selector"/>
+</shape>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_app_info_pill.xml b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_app_info_pill.xml
index 167a003..c03d240 100644
--- a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_app_info_pill.xml
+++ b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_app_info_pill.xml
@@ -19,7 +19,7 @@
     android:layout_width="@dimen/desktop_mode_handle_menu_width"
     android:layout_height="@dimen/desktop_mode_handle_menu_app_info_pill_height"
     android:orientation="horizontal"
-    android:background="@drawable/desktop_mode_decor_menu_background"
+    android:background="@drawable/desktop_mode_decor_handle_menu_background"
     android:gravity="center_vertical">
 
     <ImageView
diff --git a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_more_actions_pill.xml b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_more_actions_pill.xml
index 40a4b53..cdf4937 100644
--- a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_more_actions_pill.xml
+++ b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_more_actions_pill.xml
@@ -18,7 +18,7 @@
     android:layout_width="@dimen/desktop_mode_handle_menu_width"
     android:layout_height="@dimen/desktop_mode_handle_menu_more_actions_pill_height"
     android:orientation="vertical"
-    android:background="@drawable/desktop_mode_decor_menu_background">
+    android:background="@drawable/desktop_mode_decor_handle_menu_background">
 
     <Button
         android:id="@+id/screenshot_button"
diff --git a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_windowing_pill.xml b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_windowing_pill.xml
index 95283b9..08d9149 100644
--- a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_windowing_pill.xml
+++ b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu_windowing_pill.xml
@@ -18,7 +18,7 @@
     android:layout_width="@dimen/desktop_mode_handle_menu_width"
     android:layout_height="@dimen/desktop_mode_handle_menu_windowing_pill_height"
     android:orientation="horizontal"
-    android:background="@drawable/desktop_mode_decor_menu_background"
+    android:background="@drawable/desktop_mode_decor_handle_menu_background"
     android:gravity="center_vertical">
 
     <ImageButton
diff --git a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_maximize_menu.xml b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_maximize_menu.xml
new file mode 100644
index 0000000..0db72f7
--- /dev/null
+++ b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_maximize_menu.xml
@@ -0,0 +1,54 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2023 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    style="?android:attr/buttonBarStyle"
+    android:layout_width="@dimen/desktop_mode_maximize_menu_width"
+    android:layout_height="@dimen/desktop_mode_maximize_menu_height"
+    android:orientation="horizontal"
+    android:gravity="center"
+    android:background="@drawable/desktop_mode_maximize_menu_background">
+
+
+    <Button
+        android:id="@+id/maximize_menu_maximize_button"
+        style="?android:attr/buttonBarButtonStyle"
+        android:layout_width="120dp"
+        android:layout_height="80dp"
+        android:layout_marginRight="15dp"
+        android:color="@color/desktop_mode_maximize_menu_button"
+        android:background="@drawable/desktop_mode_maximize_menu_maximize_button_background"
+        android:stateListAnimator="@null"/>
+
+    <Button
+        android:id="@+id/maximize_menu_snap_left_button"
+        style="?android:attr/buttonBarButtonStyle"
+        android:layout_width="58dp"
+        android:layout_height="80dp"
+        android:layout_marginRight="6dp"
+        android:color="@color/desktop_mode_maximize_menu_button"
+        android:background="@drawable/desktop_mode_maximize_menu_snap_left_button_background"
+        android:stateListAnimator="@null"/>
+
+    <Button
+        android:id="@+id/maximize_menu_snap_right_button"
+        style="?android:attr/buttonBarButtonStyle"
+        android:layout_width="58dp"
+        android:layout_height="80dp"
+        android:color="@color/desktop_mode_maximize_menu_button"
+        android:background="@drawable/desktop_mode_maximize_menu_snap_right_button_background"
+        android:stateListAnimator="@null"/>
+</LinearLayout>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/res/values/colors.xml b/libs/WindowManager/Shell/res/values/colors.xml
index b2ec98b..f76a346 100644
--- a/libs/WindowManager/Shell/res/values/colors.xml
+++ b/libs/WindowManager/Shell/res/values/colors.xml
@@ -73,4 +73,8 @@
     <color name="desktop_mode_caption_menu_buttons_color_active">#00677E</color>
     <color name="desktop_mode_resize_veil_light">#EFF1F2</color>
     <color name="desktop_mode_resize_veil_dark">#1C1C17</color>
+    <color name="desktop_mode_maximize_menu_button">#DDDACD</color>
+    <color name="desktop_mode_maximize_menu_button_outline">#797869</color>
+    <color name="desktop_mode_maximize_menu_button_outline_on_hover">#606219</color>
+    <color name="desktop_mode_maximize_menu_button_on_hover">#E7E790</color>
 </resources>
diff --git a/libs/WindowManager/Shell/res/values/dimen.xml b/libs/WindowManager/Shell/res/values/dimen.xml
index 20bf81d..d0c0c02 100644
--- a/libs/WindowManager/Shell/res/values/dimen.xml
+++ b/libs/WindowManager/Shell/res/values/dimen.xml
@@ -401,6 +401,24 @@
     <!-- Height of button (32dp)  + 2 * margin (5dp each). -->
     <dimen name="freeform_decor_caption_height">42dp</dimen>
 
+    <!-- The width of the maximize menu in desktop mode. -->
+    <dimen name="desktop_mode_maximize_menu_width">287dp</dimen>
+
+    <!-- The height of the maximize menu in desktop mode. -->
+    <dimen name="desktop_mode_maximize_menu_height">112dp</dimen>
+
+    <!-- The larger of the two corner radii of the maximize menu buttons. -->
+    <dimen name="desktop_mode_maximize_menu_buttons_large_corner_radius">4dp</dimen>
+
+    <!-- The smaller of the two corner radii of the maximize menu buttons. -->
+    <dimen name="desktop_mode_maximize_menu_buttons_small_corner_radius">2dp</dimen>
+
+    <!-- The corner radius of the maximize menu. -->
+    <dimen name="desktop_mode_maximize_menu_corner_radius">8dp</dimen>
+
+    <!-- The radius of the Maximize menu shadow. -->
+    <dimen name="desktop_mode_maximize_menu_shadow_radius">8dp</dimen>
+
     <!-- The width of the handle menu in desktop mode. -->
     <dimen name="desktop_mode_handle_menu_width">216dp</dimen>
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 36d2a70..93ce91f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -196,7 +196,8 @@
             SyncTransactionQueue syncQueue,
             Transitions transitions,
             Optional<DesktopModeController> desktopModeController,
-            Optional<DesktopTasksController> desktopTasksController) {
+            Optional<DesktopTasksController> desktopTasksController,
+            RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
         if (DesktopModeStatus.isAnyEnabled()) {
             return new DesktopModeWindowDecorViewModel(
                     context,
@@ -209,7 +210,8 @@
                     syncQueue,
                     transitions,
                     desktopModeController,
-                    desktopTasksController);
+                    desktopTasksController,
+                    rootTaskDisplayAreaOrganizer);
         }
         return new CaptionWindowDecorViewModel(
                 context,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index b0f75c6..4740a9d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -105,7 +105,8 @@
 
     private val transitionAreaHeight
         get() = context.resources.getDimensionPixelSize(
-                com.android.wm.shell.R.dimen.desktop_mode_transition_area_height)
+                com.android.wm.shell.R.dimen.desktop_mode_transition_area_height
+        )
 
     // This is public to avoid cyclic dependency; it is set by SplitScreenController
     lateinit var splitScreenController: SplitScreenController
@@ -485,6 +486,55 @@
         }
     }
 
+    /**
+     * Quick-resize to the right or left half of the stable bounds.
+     *
+     * @param position the portion of the screen (RIGHT or LEFT) we want to snap the task to.
+     */
+    fun snapToHalfScreen(
+            taskInfo: RunningTaskInfo,
+            windowDecor: DesktopModeWindowDecoration,
+            position: SnapPosition
+    ) {
+        val displayLayout = displayController.getDisplayLayout(taskInfo.displayId) ?: return
+
+        val stableBounds = Rect()
+        displayLayout.getStableBounds(stableBounds)
+
+        val destinationWidth = stableBounds.width() / 2
+        val destinationBounds = when (position) {
+            SnapPosition.LEFT -> {
+                Rect(
+                        stableBounds.left,
+                        stableBounds.top,
+                        stableBounds.left + destinationWidth,
+                        stableBounds.bottom
+                )
+            }
+            SnapPosition.RIGHT -> {
+                Rect(
+                        stableBounds.right - destinationWidth,
+                        stableBounds.top,
+                        stableBounds.right,
+                        stableBounds.bottom
+                )
+            }
+        }
+
+        if (destinationBounds == taskInfo.configuration.windowConfiguration.bounds) return
+
+        val wct = WindowContainerTransaction().setBounds(taskInfo.token, destinationBounds)
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            toggleResizeDesktopTaskTransitionHandler.startTransition(
+                    wct,
+                    taskInfo.taskId,
+                    windowDecor
+            )
+        } else {
+            shellTaskOrganizer.applyTransaction(wct)
+        }
+    }
+
     private fun getDefaultDesktopTaskBounds(density: Float, stableBounds: Rect, outBounds: Rect) {
         val width = (DESKTOP_MODE_DEFAULT_WIDTH_DP * density + 0.5f).toInt()
         val height = (DESKTOP_MODE_DEFAULT_HEIGHT_DP * density + 0.5f).toInt()
@@ -1077,4 +1127,7 @@
             return DESKTOP_DENSITY_OVERRIDE in DESKTOP_DENSITY_ALLOWED_RANGE
         }
     }
+
+    /** The positions on a screen that a task can snap to. */
+    enum class SnapPosition { RIGHT, LEFT }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index 29fff03..026e973 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -63,6 +63,7 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.wm.shell.R;
+import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayLayout;
@@ -70,6 +71,7 @@
 import com.android.wm.shell.desktopmode.DesktopModeController;
 import com.android.wm.shell.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
+import com.android.wm.shell.desktopmode.DesktopTasksController.SnapPosition;
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter;
 import com.android.wm.shell.splitscreen.SplitScreen;
 import com.android.wm.shell.splitscreen.SplitScreenController;
@@ -120,6 +122,7 @@
     private MoveToDesktopAnimator mMoveToDesktopAnimator;
     private final Rect mDragToDesktopAnimationStartBounds = new Rect();
     private final DesktopModeKeyguardChangeListener mDesktopModeKeyguardChangeListener;
+    private final RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
 
     public DesktopModeWindowDecorViewModel(
             Context context,
@@ -132,7 +135,8 @@
             SyncTransactionQueue syncQueue,
             Transitions transitions,
             Optional<DesktopModeController> desktopModeController,
-            Optional<DesktopTasksController> desktopTasksController
+            Optional<DesktopTasksController> desktopTasksController,
+            RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer
     ) {
         this(
                 context,
@@ -149,7 +153,8 @@
                 new DesktopModeWindowDecoration.Factory(),
                 new InputMonitorFactory(),
                 SurfaceControl.Transaction::new,
-                new DesktopModeKeyguardChangeListener());
+                new DesktopModeKeyguardChangeListener(),
+                rootTaskDisplayAreaOrganizer);
     }
 
     @VisibleForTesting
@@ -168,7 +173,8 @@
             DesktopModeWindowDecoration.Factory desktopModeWindowDecorFactory,
             InputMonitorFactory inputMonitorFactory,
             Supplier<SurfaceControl.Transaction> transactionFactory,
-            DesktopModeKeyguardChangeListener desktopModeKeyguardChangeListener) {
+            DesktopModeKeyguardChangeListener desktopModeKeyguardChangeListener,
+            RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
         mContext = context;
         mMainHandler = mainHandler;
         mMainChoreographer = mainChoreographer;
@@ -185,6 +191,7 @@
         mInputMonitorFactory = inputMonitorFactory;
         mTransactionFactory = transactionFactory;
         mDesktopModeKeyguardChangeListener = desktopModeKeyguardChangeListener;
+        mRootTaskDisplayAreaOrganizer = rootTaskDisplayAreaOrganizer;
 
         shellInit.addInitCallback(this::onInit, this);
     }
@@ -318,7 +325,8 @@
     }
 
     private class DesktopModeTouchEventListener extends GestureDetector.SimpleOnGestureListener
-            implements View.OnClickListener, View.OnTouchListener, DragDetector.MotionEventHandler {
+            implements View.OnClickListener, View.OnTouchListener, View.OnLongClickListener,
+            DragDetector.MotionEventHandler{
 
         private final int mTaskId;
         private final WindowContainerToken mTaskToken;
@@ -355,9 +363,11 @@
                             .getTaskInfo(remainingTaskPosition);
                     mSplitScreenController.moveTaskToFullscreen(remainingTask.taskId);
                 }
+                decoration.closeMaximizeMenu();
             } else if (id == R.id.back_button) {
                 mTaskOperations.injectBackKey();
             } else if (id == R.id.caption_handle || id == R.id.open_menu_button) {
+                decoration.closeMaximizeMenu();
                 if (!decoration.isHandleMenuActive()) {
                     moveTaskToFront(mTaskOrganizer.getRunningTaskInfo(mTaskId));
                     decoration.createHandleMenu();
@@ -391,13 +401,35 @@
                     // TODO(b/278084491): dev option to enable display switching
                     //  remove when select is implemented
                     mDesktopTasksController.ifPresent(c -> c.moveToNextDisplay(mTaskId));
-                    decoration.closeHandleMenu();
                 }
             } else if (id == R.id.maximize_window) {
+                moveTaskToFront(decoration.mTaskInfo);
+                if (decoration.isMaximizeMenuActive()) {
+                    decoration.closeMaximizeMenu();
+                    return;
+                }
                 final RunningTaskInfo taskInfo = decoration.mTaskInfo;
                 mDesktopTasksController.ifPresent(c -> c.toggleDesktopTaskSize(
                         taskInfo, decoration));
                 decoration.closeHandleMenu();
+            } else if (id == R.id.maximize_menu_maximize_button) {
+                final RunningTaskInfo taskInfo = decoration.mTaskInfo;
+                mDesktopTasksController.ifPresent(c -> c.toggleDesktopTaskSize(
+                        taskInfo, mWindowDecorByTaskId.get(taskInfo.taskId)));
+                decoration.closeHandleMenu();
+                decoration.closeMaximizeMenu();
+            } else if (id == R.id.maximize_menu_snap_left_button) {
+                final RunningTaskInfo taskInfo = decoration.mTaskInfo;
+                mDesktopTasksController.ifPresent(c -> c.snapToHalfScreen(
+                        taskInfo, mWindowDecorByTaskId.get(taskInfo.taskId), SnapPosition.LEFT));
+                decoration.closeHandleMenu();
+                decoration.closeMaximizeMenu();
+            } else if (id == R.id.maximize_menu_snap_right_button) {
+                final RunningTaskInfo taskInfo = decoration.mTaskInfo;
+                mDesktopTasksController.ifPresent(c -> c.snapToHalfScreen(
+                        taskInfo, mWindowDecorByTaskId.get(taskInfo.taskId), SnapPosition.RIGHT));
+                decoration.closeHandleMenu();
+                decoration.closeMaximizeMenu();
             }
         }
 
@@ -412,6 +444,23 @@
             return mDragDetector.onMotionEvent(v, e);
         }
 
+        @Override
+        public boolean onLongClick(View v) {
+            final int id = v.getId();
+            if (id == R.id.maximize_window) {
+                final DesktopModeWindowDecoration decoration = mWindowDecorByTaskId.get(mTaskId);
+                moveTaskToFront(decoration.mTaskInfo);
+                if (decoration.isMaximizeMenuActive()) {
+                    decoration.closeMaximizeMenu();
+                } else {
+                    decoration.closeHandleMenu();
+                    decoration.createMaximizeMenu();
+                }
+                return true;
+            }
+            return false;
+        }
+
         private void moveTaskToFront(RunningTaskInfo taskInfo) {
             if (!taskInfo.isFocused) {
                 mDesktopTasksController.ifPresent(c -> c.moveTaskToFront(taskInfo));
@@ -875,7 +924,8 @@
                         taskSurface,
                         mMainHandler,
                         mMainChoreographer,
-                        mSyncQueue);
+                        mSyncQueue,
+                        mRootTaskDisplayAreaOrganizer);
         mWindowDecorByTaskId.put(taskInfo.taskId, windowDecoration);
         windowDecoration.createResizeVeil();
 
@@ -884,7 +934,8 @@
         final DesktopModeTouchEventListener touchEventListener =
                 new DesktopModeTouchEventListener(taskInfo, dragPositioningCallback);
 
-        windowDecoration.setCaptionListeners(touchEventListener, touchEventListener);
+        windowDecoration.setCaptionListeners(
+                touchEventListener, touchEventListener, touchEventListener);
         windowDecoration.setCornersListener(mCornersListener);
         windowDecoration.setDragPositioningCallback(dragPositioningCallback);
         windowDecoration.setDragDetector(touchEventListener.mDragDetector);
@@ -911,7 +962,9 @@
             implements DragPositioningCallbackUtility.DragStartListener {
         @Override
         public void onDragStart(int taskId) {
-            mWindowDecorByTaskId.get(taskId).closeHandleMenu();
+            final DesktopModeWindowDecoration decoration = mWindowDecorByTaskId.get(taskId);
+            decoration.closeHandleMenu();
+            decoration.closeMaximizeMenu();
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
index a359395..a75dce2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
@@ -23,6 +23,7 @@
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.content.res.Configuration;
+import android.content.res.Resources;
 import android.graphics.Point;
 import android.graphics.PointF;
 import android.graphics.Rect;
@@ -36,13 +37,16 @@
 import android.view.SurfaceControl;
 import android.view.View;
 import android.view.ViewConfiguration;
+import android.widget.ImageButton;
 import android.window.WindowContainerTransaction;
 
 import com.android.internal.policy.ScreenDecorationsUtils;
 import com.android.launcher3.icons.IconProvider;
 import com.android.wm.shell.R;
+import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
+import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
@@ -69,6 +73,7 @@
     private DesktopModeWindowDecorationViewHolder mWindowDecorViewHolder;
     private View.OnClickListener mOnCaptionButtonClickListener;
     private View.OnTouchListener mOnCaptionTouchListener;
+    private View.OnLongClickListener mOnCaptionLongClickListener;
     private DragPositioningCallback mDragPositioningCallback;
     private DragResizeInputListener mDragResizeListener;
     private DragDetector mDragDetector;
@@ -80,6 +85,8 @@
     private final Point mPositionInParent = new Point();
     private HandleMenu mHandleMenu;
 
+    private MaximizeMenu mMaximizeMenu;
+
     private ResizeVeil mResizeVeil;
 
     private Drawable mAppIcon;
@@ -89,6 +96,7 @@
 
     private final Set<IBinder> mTransitionsPausingRelayout = new HashSet<>();
     private int mRelayoutBlock;
+    private final RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
 
     DesktopModeWindowDecoration(
             Context context,
@@ -98,12 +106,14 @@
             SurfaceControl taskSurface,
             Handler handler,
             Choreographer choreographer,
-            SyncTransactionQueue syncQueue) {
+            SyncTransactionQueue syncQueue,
+            RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
         super(context, displayController, taskOrganizer, taskInfo, taskSurface);
 
         mHandler = handler;
         mChoreographer = choreographer;
         mSyncQueue = syncQueue;
+        mRootTaskDisplayAreaOrganizer = rootTaskDisplayAreaOrganizer;
 
         loadAppInfo();
     }
@@ -121,9 +131,11 @@
 
     void setCaptionListeners(
             View.OnClickListener onCaptionButtonClickListener,
-            View.OnTouchListener onCaptionTouchListener) {
+            View.OnTouchListener onCaptionTouchListener,
+            View.OnLongClickListener onLongClickListener) {
         mOnCaptionButtonClickListener = onCaptionButtonClickListener;
         mOnCaptionTouchListener = onCaptionTouchListener;
+        mOnCaptionLongClickListener = onLongClickListener;
     }
 
     void setCornersListener(TaskCornersListener cornersListener) {
@@ -207,6 +219,7 @@
                         mResult.mRootView,
                         mOnCaptionTouchListener,
                         mOnCaptionButtonClickListener,
+                        mOnCaptionLongClickListener,
                         mAppName,
                         mAppIcon
                 );
@@ -218,6 +231,7 @@
 
         if (!mTaskInfo.isFocused) {
             closeHandleMenu();
+            closeMaximizeMenu();
         }
 
         if (!isDragResizeable) {
@@ -255,6 +269,52 @@
             mCornersListener.onTaskCornersChanged(mTaskInfo.taskId, getGlobalCornersRegion());
         }
         mPositionInParent.set(mTaskInfo.positionInParent);
+
+        if (isMaximizeMenuActive()) {
+            if (!mTaskInfo.isVisible()) {
+                closeMaximizeMenu();
+            } else {
+                mMaximizeMenu.positionMenu(calculateMaximizeMenuPosition(), startT);
+            }
+        }
+    }
+
+    private PointF calculateMaximizeMenuPosition() {
+        final PointF position = new PointF();
+        final Resources resources = mContext.getResources();
+        final DisplayLayout displayLayout =
+                mDisplayController.getDisplayLayout(mTaskInfo.displayId);
+        if (displayLayout == null) return position;
+
+        final int displayWidth = displayLayout.width();
+        final int displayHeight = displayLayout.height();
+        final int captionHeight = loadDimensionPixelSize(
+                resources, R.dimen.freeform_decor_caption_height);
+
+        final ImageButton maximizeWindowButton =
+                mResult.mRootView.findViewById(R.id.maximize_window);
+        final int[] maximizeButtonLocation = new int[2];
+        maximizeWindowButton.getLocationInWindow(maximizeButtonLocation);
+
+        final int menuWidth = loadDimensionPixelSize(
+                resources, R.dimen.desktop_mode_maximize_menu_width);
+        final int menuHeight = loadDimensionPixelSize(
+                resources, R.dimen.desktop_mode_maximize_menu_height);
+
+        float menuLeft = (mPositionInParent.x + maximizeButtonLocation[0]);
+        float menuTop = (mPositionInParent.y + captionHeight);
+        final float menuRight = menuLeft + menuWidth;
+        final float menuBottom = menuTop + menuHeight;
+
+        // If the menu is out of screen bounds, shift it up/left as needed
+        if (menuRight > displayWidth) {
+            menuLeft = (displayWidth - menuWidth);
+        }
+        if (menuBottom > displayHeight) {
+            menuTop = (displayHeight - menuHeight);
+        }
+
+        return new PointF(menuLeft, menuTop);
     }
 
     boolean isHandleMenuActive() {
@@ -335,6 +395,29 @@
     }
 
     /**
+     * Create and display maximize menu window
+     */
+    void createMaximizeMenu() {
+        mMaximizeMenu = new MaximizeMenu(mSyncQueue, mRootTaskDisplayAreaOrganizer,
+                mDisplayController, mTaskInfo, mOnCaptionButtonClickListener, mContext,
+                calculateMaximizeMenuPosition(), mSurfaceControlTransactionSupplier);
+        mMaximizeMenu.show();
+    }
+
+    /**
+     * Close the maximize menu window
+     */
+    void closeMaximizeMenu() {
+        if (!isMaximizeMenuActive()) return;
+        mMaximizeMenu.close();
+        mMaximizeMenu = null;
+    }
+
+    boolean isMaximizeMenuActive() {
+        return mMaximizeMenu != null;
+    }
+
+    /**
      * Create and display handle menu window
      */
     void createHandleMenu() {
@@ -532,7 +615,8 @@
                 SurfaceControl taskSurface,
                 Handler handler,
                 Choreographer choreographer,
-                SyncTransactionQueue syncQueue) {
+                SyncTransactionQueue syncQueue,
+                RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
             return new DesktopModeWindowDecoration(
                     context,
                     displayController,
@@ -541,7 +625,8 @@
                     taskSurface,
                     handler,
                     choreographer,
-                    syncQueue);
+                    syncQueue,
+                    rootTaskDisplayAreaOrganizer);
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeMenu.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeMenu.kt
new file mode 100644
index 0000000..4dc98e4
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeMenu.kt
@@ -0,0 +1,154 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.windowdecor
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.content.Context
+import android.content.res.Resources
+import android.graphics.PixelFormat
+import android.graphics.PointF
+import android.view.LayoutInflater
+import android.view.SurfaceControl
+import android.view.SurfaceControl.Transaction
+import android.view.SurfaceControlViewHost
+import android.view.View.OnClickListener
+import android.view.WindowManager
+import android.view.WindowlessWindowManager
+import android.widget.Button
+import android.window.TaskConstants
+import com.android.wm.shell.R
+import com.android.wm.shell.RootTaskDisplayAreaOrganizer
+import com.android.wm.shell.common.DisplayController
+import com.android.wm.shell.common.SyncTransactionQueue
+import com.android.wm.shell.windowdecor.WindowDecoration.AdditionalWindow
+import java.util.function.Supplier
+
+
+/**
+ *  Menu that appears when user long clicks the maximize button. Gives the user the option to
+ *  maximize the task or snap the task to the right or left half of the screen.
+ */
+class MaximizeMenu(
+        private val syncQueue: SyncTransactionQueue,
+        private val rootTdaOrganizer: RootTaskDisplayAreaOrganizer,
+        private val displayController: DisplayController,
+        private val taskInfo: RunningTaskInfo,
+        private val onClickListener: OnClickListener,
+        private val decorWindowContext: Context,
+        private val menuPosition: PointF,
+        private val transactionSupplier: Supplier<Transaction> = Supplier { Transaction() }
+) {
+    private var maximizeMenu: AdditionalWindow? = null
+    private lateinit var viewHost: SurfaceControlViewHost
+    private lateinit var leash: SurfaceControl
+    private val shadowRadius = loadDimensionPixelSize(
+            R.dimen.desktop_mode_maximize_menu_shadow_radius
+    ).toFloat()
+    private val cornerRadius = loadDimensionPixelSize(
+            R.dimen.desktop_mode_maximize_menu_corner_radius
+    ).toFloat()
+
+    /** Position the menu relative to the caption's position. */
+    fun positionMenu(position: PointF, t: Transaction) {
+        menuPosition.set(position)
+        t.setPosition(leash, menuPosition.x, menuPosition.y)
+    }
+
+    /** Creates and shows the maximize window. */
+    fun show() {
+        if (maximizeMenu != null) return
+        createMaximizeMenu()
+        setupMaximizeMenu()
+    }
+
+    /** Closes the maximize window and releases its view. */
+    fun close() {
+        maximizeMenu?.releaseView()
+        maximizeMenu = null
+    }
+
+    /** Create a maximize menu that is attached to the display area. */
+    private fun createMaximizeMenu() {
+        val t = transactionSupplier.get()
+        val v = LayoutInflater.from(decorWindowContext).inflate(
+                R.layout.desktop_mode_window_decor_maximize_menu,
+                null // Root
+        )
+        val builder = SurfaceControl.Builder()
+        rootTdaOrganizer.attachToDisplayArea(taskInfo.displayId, builder)
+        leash = builder
+                .setName("Maximize Menu")
+                .setContainerLayer()
+                .build()
+        val menuWidth = loadDimensionPixelSize(R.dimen.desktop_mode_maximize_menu_width)
+        val menuHeight = loadDimensionPixelSize(R.dimen.desktop_mode_maximize_menu_height)
+        val lp = WindowManager.LayoutParams(
+                menuWidth,
+                menuHeight,
+                WindowManager.LayoutParams.TYPE_APPLICATION,
+                WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE,
+                PixelFormat.TRANSPARENT
+        )
+        lp.title = "Maximize Menu for Task=" + taskInfo.taskId
+        lp.setTrustedOverlay()
+        val windowManager = WindowlessWindowManager(
+                taskInfo.configuration,
+                leash,
+                null // HostInputToken
+        )
+        viewHost = SurfaceControlViewHost(decorWindowContext,
+                displayController.getDisplay(taskInfo.displayId), windowManager,
+                "MaximizeMenu")
+        viewHost.setView(v, lp)
+
+        // Bring menu to front when open
+        t.setLayer(leash, TaskConstants.TASK_CHILD_LAYER_FLOATING_MENU)
+                .setPosition(leash, menuPosition.x, menuPosition.y)
+                .setWindowCrop(leash, menuWidth, menuHeight)
+                .setShadowRadius(leash, shadowRadius)
+                .setCornerRadius(leash, cornerRadius)
+                .show(leash)
+        maximizeMenu = AdditionalWindow(leash, viewHost, transactionSupplier)
+
+        syncQueue.runInSync { transaction ->
+            transaction.merge(t)
+            t.close()
+        }
+    }
+
+    private fun loadDimensionPixelSize(resourceId: Int): Int {
+        return if (resourceId == Resources.ID_NULL) {
+            0
+        } else {
+            decorWindowContext.resources.getDimensionPixelSize(resourceId)
+        }
+    }
+
+    private fun setupMaximizeMenu() {
+        val maximizeMenuView = maximizeMenu?.mWindowViewHost?.view ?: return
+
+        maximizeMenuView.findViewById<Button>(
+                R.id.maximize_menu_maximize_button
+        ).setOnClickListener(onClickListener)
+        maximizeMenuView.findViewById<Button>(
+                R.id.maximize_menu_snap_right_button
+        ).setOnClickListener(onClickListener)
+        maximizeMenuView.findViewById<Button>(
+                R.id.maximize_menu_snap_left_button
+        ).setOnClickListener(onClickListener)
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/DesktopModeAppControlsWindowDecorationViewHolder.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/DesktopModeAppControlsWindowDecorationViewHolder.kt
index a9eb882..6b59cce 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/DesktopModeAppControlsWindowDecorationViewHolder.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/DesktopModeAppControlsWindowDecorationViewHolder.kt
@@ -5,6 +5,7 @@
 import android.graphics.drawable.Drawable
 import android.graphics.drawable.GradientDrawable
 import android.view.View
+import android.view.View.OnLongClickListener
 import android.widget.ImageButton
 import android.widget.ImageView
 import android.widget.TextView
@@ -19,6 +20,7 @@
         rootView: View,
         onCaptionTouchListener: View.OnTouchListener,
         onCaptionButtonClickListener: View.OnClickListener,
+        onLongClickListener: OnLongClickListener,
         appName: CharSequence,
         appIcon: Drawable
 ) : DesktopModeWindowDecorationViewHolder(rootView) {
@@ -39,6 +41,7 @@
         openMenuButton.setOnTouchListener(onCaptionTouchListener)
         closeWindowButton.setOnClickListener(onCaptionButtonClickListener)
         maximizeWindowButton.setOnClickListener(onCaptionButtonClickListener)
+        maximizeWindowButton.onLongClickListener = onLongClickListener
         closeWindowButton.setOnTouchListener(onCaptionTouchListener)
         appNameTextView.text = appName
         appIconImageView.setImageDrawable(appIcon)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.java
index 596d6dd..7f0465a 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.java
@@ -48,6 +48,7 @@
 
 import androidx.test.filters.SmallTest;
 
+import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestRunningTaskInfoBuilder;
@@ -100,6 +101,7 @@
     @Mock private ShellInit mShellInit;
     @Mock private DesktopModeWindowDecorViewModel.DesktopModeKeyguardChangeListener
             mDesktopModeKeyguardChangeListener;
+    @Mock private RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
     private final List<InputManager> mMockInputManagers = new ArrayList<>();
 
     private DesktopModeWindowDecorViewModel mDesktopModeWindowDecorViewModel;
@@ -109,27 +111,28 @@
         mMockInputManagers.add(mInputManager);
 
         mDesktopModeWindowDecorViewModel =
-            new DesktopModeWindowDecorViewModel(
-                mContext,
-                mMainHandler,
-                mMainChoreographer,
-                mShellInit,
-                mTaskOrganizer,
-                mDisplayController,
-                mShellController,
-                mSyncQueue,
-                mTransitions,
-                Optional.of(mDesktopModeController),
-                Optional.of(mDesktopTasksController),
-                mDesktopModeWindowDecorFactory,
-                mMockInputMonitorFactory,
-                mTransactionFactory,
-                mDesktopModeKeyguardChangeListener
-            );
+                new DesktopModeWindowDecorViewModel(
+                        mContext,
+                        mMainHandler,
+                        mMainChoreographer,
+                        mShellInit,
+                        mTaskOrganizer,
+                        mDisplayController,
+                        mShellController,
+                        mSyncQueue,
+                        mTransitions,
+                        Optional.of(mDesktopModeController),
+                        Optional.of(mDesktopTasksController),
+                        mDesktopModeWindowDecorFactory,
+                        mMockInputMonitorFactory,
+                        mTransactionFactory,
+                        mDesktopModeKeyguardChangeListener,
+                        mRootTaskDisplayAreaOrganizer
+                );
 
         doReturn(mDesktopModeWindowDecoration)
-            .when(mDesktopModeWindowDecorFactory)
-            .create(any(), any(), any(), any(), any(), any(), any(), any());
+                .when(mDesktopModeWindowDecorFactory)
+                .create(any(), any(), any(), any(), any(), any(), any(), any(), any());
         doReturn(mTransaction).when(mTransactionFactory).get();
         doReturn(mDisplayLayout).when(mDisplayController).getDisplayLayout(anyInt());
         doReturn(STABLE_INSETS).when(mDisplayLayout).stableInsets();
@@ -172,7 +175,8 @@
                         surfaceControl,
                         mMainHandler,
                         mMainChoreographer,
-                        mSyncQueue);
+                        mSyncQueue,
+                        mRootTaskDisplayAreaOrganizer);
         verify(mDesktopModeWindowDecoration).close();
     }
 
@@ -205,7 +209,8 @@
                         surfaceControl,
                         mMainHandler,
                         mMainChoreographer,
-                        mSyncQueue);
+                        mSyncQueue,
+                        mRootTaskDisplayAreaOrganizer);
     }
 
     @Test
@@ -291,7 +296,7 @@
                     taskInfo, surfaceControl, startT, finishT);
         });
         verify(mDesktopModeWindowDecorFactory, never())
-                .create(any(), any(), any(), any(), any(), any(), any(), any());
+                .create(any(), any(), any(), any(), any(), any(), any(), any(), any());
     }
 
     private void runOnMainThread(Runnable r) throws Exception {
@@ -307,10 +312,10 @@
     private static ActivityManager.RunningTaskInfo createTaskInfo(int taskId,
             int displayId, @WindowConfiguration.WindowingMode int windowingMode) {
         ActivityManager.RunningTaskInfo taskInfo =
-                 new TestRunningTaskInfoBuilder()
-                .setDisplayId(displayId)
-                .setVisible(true)
-                .build();
+                new TestRunningTaskInfoBuilder()
+                        .setDisplayId(displayId)
+                        .setVisible(true)
+                        .build();
         taskInfo.taskId = taskId;
         taskInfo.configuration.windowConfiguration.setWindowingMode(windowingMode);
         return taskInfo;
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index b472982..2077af8 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -135,7 +135,10 @@
 aconfig_declarations {
     name: "systemui_aconfig_flags",
     package: "com.android.systemui.aconfig",
-    srcs: ["src/com/android/systemui/aconfig/systemui.aconfig"],
+    srcs: [
+        "src/com/android/systemui/aconfig/systemui.aconfig",
+        "src/com/android/systemui/accessibility/aconfig/accessibility.aconfig",
+    ],
 }
 
 java_aconfig_library {
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/aconfig/accessibility.aconfig b/packages/SystemUI/src/com/android/systemui/accessibility/aconfig/accessibility.aconfig
new file mode 100644
index 0000000..91c5551
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/aconfig/accessibility.aconfig
@@ -0,0 +1,7 @@
+package: "com.android.systemui.aconfig"
+flag {
+    name: "floating_menu_overlaps_nav_bars_flag"
+    namespace: "accessibility"
+    description: "Adjusts bounds to allow the floating menu to render on top of navigation bars."
+    bug: "283768342"
+}
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuView.java b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuView.java
index 47770fa..f29077d 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuView.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuView.java
@@ -38,6 +38,7 @@
 import androidx.recyclerview.widget.RecyclerViewAccessibilityDelegate;
 
 import com.android.internal.accessibility.dialog.AccessibilityTarget;
+import com.android.systemui.aconfig.Flags;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -284,6 +285,22 @@
     void updateMenuMoveToTucked(boolean isMoveToTucked) {
         mIsMoveToTucked = isMoveToTucked;
         mMenuViewModel.updateMenuMoveToTucked(isMoveToTucked);
+
+        if (Flags.floatingMenuOverlapsNavBarsFlag()) {
+            if (isMoveToTucked) {
+                final float halfWidth = getMenuWidth() / 2.0f;
+                final boolean isOnLeftSide = mMenuAnimationController.isOnLeftSide();
+                final Rect clipBounds = new Rect(
+                        (int) (!isOnLeftSide ? 0 : halfWidth),
+                        0,
+                        (int) (!isOnLeftSide ? halfWidth : getMenuWidth()),
+                        getMenuHeight()
+                );
+                setClipBounds(clipBounds);
+            } else {
+                setClipBounds(null);
+            }
+        }
     }
 
 
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewAppearance.java b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewAppearance.java
index 3cd250f..3822936 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewAppearance.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewAppearance.java
@@ -35,6 +35,7 @@
 import androidx.annotation.DimenRes;
 
 import com.android.systemui.R;
+import com.android.systemui.aconfig.Flags;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -154,8 +155,10 @@
         final int margin = getMenuMargin();
         final Rect draggableBounds = new Rect(getWindowAvailableBounds());
 
-        // Initializes start position for mapping the translation of the menu view.
-        draggableBounds.offsetTo(/* newLeft= */ 0, /* newTop= */ 0);
+        if (!Flags.floatingMenuOverlapsNavBarsFlag()) {
+            // Initializes start position for mapping the translation of the menu view.
+            draggableBounds.offsetTo(/* newLeft= */ 0, /* newTop= */ 0);
+        }
 
         draggableBounds.top += margin;
         draggableBounds.right -= getMenuWidth();
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayerController.java b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayerController.java
index c52ecc5..cc18c30 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayerController.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayerController.java
@@ -24,6 +24,7 @@
 import android.view.WindowManager;
 import android.view.accessibility.AccessibilityManager;
 
+import com.android.systemui.aconfig.Flags;
 import com.android.systemui.util.settings.SecureSettings;
 
 /**
@@ -77,8 +78,15 @@
         params.receiveInsetsIgnoringZOrder = true;
         params.privateFlags |= PRIVATE_FLAG_EXCLUDE_FROM_SCREEN_MAGNIFICATION;
         params.windowAnimations = android.R.style.Animation_Translucent;
-        params.setFitInsetsTypes(
-                WindowInsets.Type.systemBars() | WindowInsets.Type.displayCutout());
+        // Insets are configured to allow the menu to display over navigation and system bars.
+        if (Flags.floatingMenuOverlapsNavBarsFlag()) {
+            params.setFitInsetsTypes(0);
+            params.layoutInDisplayCutoutMode =
+                    WindowManager.LayoutParams.LAYOUT_IN_DISPLAY_CUTOUT_MODE_ALWAYS;
+        } else {
+            params.setFitInsetsTypes(
+                    WindowInsets.Type.systemBars() | WindowInsets.Type.displayCutout());
+        }
         return params;
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
index c0d807a..98f2fee 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
@@ -28,8 +28,18 @@
 import java.util.function.Consumer
 import javax.inject.Inject
 
+/** Processes a screenshot request sent from [ScreenshotHelper]. */
+interface ScreenshotRequestProcessor {
+    /**
+     * Inspects the incoming ScreenshotData, potentially modifying it based upon policy.
+     *
+     * @param screenshot the screenshot to process
+     */
+    suspend fun process(screenshot: ScreenshotData): ScreenshotData
+}
+
 /**
- * Processes a screenshot request sent from {@link ScreenshotHelper}.
+ * Implementation of [ScreenshotRequestProcessor]
  */
 @SysUISingleton
 class RequestProcessor @Inject constructor(
@@ -38,7 +48,7 @@
         private val flags: FeatureFlags,
         /** For the Java Async version, to invoke the callback. */
         @Application private val mainScope: CoroutineScope
-) {
+) : ScreenshotRequestProcessor {
     /**
      * Inspects the incoming request, returning a potentially modified request depending on policy.
      *
@@ -57,7 +67,6 @@
         // regardless of the managed profile status.
 
         if (request.type != TAKE_SCREENSHOT_PROVIDED_IMAGE) {
-
             val info = policy.findPrimaryContent(policy.getDefaultDisplayId())
             Log.d(TAG, "findPrimaryContent: $info")
 
@@ -99,12 +108,7 @@
         }
     }
 
-    /**
-     * Inspects the incoming ScreenshotData, potentially modifying it based upon policy.
-     *
-     * @param screenshot the screenshot to process
-     */
-    suspend fun process(screenshot: ScreenshotData): ScreenshotData {
+    override suspend fun process(screenshot: ScreenshotData): ScreenshotData {
         var result = screenshot
 
         // Apply work profile screenshots policy:
@@ -116,7 +120,7 @@
         // regardless of the managed profile status.
 
         if (screenshot.type != TAKE_SCREENSHOT_PROVIDED_IMAGE) {
-            val info = policy.findPrimaryContent(policy.getDefaultDisplayId())
+            val info = policy.findPrimaryContent(screenshot.displayId)
             Log.d(TAG, "findPrimaryContent: $info")
             result.taskId = info.taskId
             result.topComponent = info.component
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index b59106e..cf782b7 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -100,11 +100,14 @@
 import com.android.systemui.flags.Flags;
 import com.android.systemui.screenshot.ScreenshotController.SavedImageData.ActionTransition;
 import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback;
-import com.android.systemui.settings.DisplayTracker;
 import com.android.systemui.util.Assert;
 
 import com.google.common.util.concurrent.ListenableFuture;
 
+import dagger.assisted.Assisted;
+import dagger.assisted.AssistedFactory;
+import dagger.assisted.AssistedInject;
+
 import java.io.File;
 import java.util.List;
 import java.util.concurrent.CancellationException;
@@ -118,7 +121,6 @@
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
-import javax.inject.Inject;
 
 /**
  * Controls the state and flow for screenshots.
@@ -275,7 +277,7 @@
     private final ScrollCaptureClient mScrollCaptureClient;
     private final PhoneWindow mWindow;
     private final DisplayManager mDisplayManager;
-    private final DisplayTracker mDisplayTracker;
+    private final int mDisplayId;
     private final ScrollCaptureController mScrollCaptureController;
     private final LongScreenshotData mLongScreenshotHolder;
     private final boolean mIsLowRamDevice;
@@ -314,7 +316,8 @@
                     | ActivityInfo.CONFIG_SCREEN_LAYOUT
                     | ActivityInfo.CONFIG_ASSETS_PATHS);
 
-    @Inject
+
+    @AssistedInject
     ScreenshotController(
             Context context,
             FeatureFlags flags,
@@ -335,7 +338,7 @@
             UserManager userManager,
             AssistContentRequester assistContentRequester,
             MessageContainerController messageContainerController,
-            DisplayTracker displayTracker
+            @Assisted int displayId
     ) {
         mScreenshotSmartActions = screenshotSmartActions;
         mNotificationsController = screenshotNotificationsController;
@@ -360,9 +363,9 @@
             dismissScreenshot(SCREENSHOT_INTERACTION_TIMEOUT);
         });
 
+        mDisplayId = displayId;
         mDisplayManager = requireNonNull(context.getSystemService(DisplayManager.class));
-        mDisplayTracker = displayTracker;
-        final Context displayContext = context.createDisplayContext(getDefaultDisplay());
+        final Context displayContext = context.createDisplayContext(getDisplay());
         mContext = (WindowContext) displayContext.createWindowContext(TYPE_SCREENSHOT, null);
         mWindowManager = mContext.getSystemService(WindowManager.class);
         mFlags = flags;
@@ -406,7 +409,7 @@
         if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_FULLSCREEN) {
             Rect bounds = getFullScreenRect();
             screenshot.setBitmap(
-                    mImageCapture.captureDisplay(mDisplayTracker.getDefaultDisplayId(), bounds));
+                    mImageCapture.captureDisplay(mDisplayId, bounds));
             screenshot.setScreenBounds(bounds);
         }
 
@@ -638,7 +641,7 @@
                 setWindowFocusable(false);
             }
         }, mActionExecutor, mFlags);
-        mScreenshotView.setDefaultDisplay(mDisplayTracker.getDefaultDisplayId());
+        mScreenshotView.setDefaultDisplay(mDisplayId);
         mScreenshotView.setDefaultTimeoutMillis(mScreenshotHandler.getDefaultTimeoutMillis());
 
         mScreenshotView.setOnKeyListener((v, keyCode, event) -> {
@@ -727,8 +730,8 @@
         if (mLastScrollCaptureRequest != null) {
             mLastScrollCaptureRequest.cancel(true);
         }
-        final ListenableFuture<ScrollCaptureResponse> future =
-                mScrollCaptureClient.request(mDisplayTracker.getDefaultDisplayId());
+        final ListenableFuture<ScrollCaptureResponse> future = mScrollCaptureClient.request(
+                mDisplayId);
         mLastScrollCaptureRequest = future;
         mLastScrollCaptureRequest.addListener(() ->
                 onScrollCaptureResponseReady(future, owner), mMainExecutor);
@@ -758,9 +761,8 @@
             final ScrollCaptureResponse response = mLastScrollCaptureResponse;
             mScreenshotView.showScrollChip(response.getPackageName(), /* onClick */ () -> {
                 DisplayMetrics displayMetrics = new DisplayMetrics();
-                getDefaultDisplay().getRealMetrics(displayMetrics);
-                Bitmap newScreenshot = mImageCapture.captureDisplay(
-                        mDisplayTracker.getDefaultDisplayId(),
+                getDisplay().getRealMetrics(displayMetrics);
+                Bitmap newScreenshot = mImageCapture.captureDisplay(mDisplayId,
                         new Rect(0, 0, displayMetrics.widthPixels, displayMetrics.heightPixels));
 
                 mScreenshotView.prepareScrollingTransition(response, mScreenBitmap, newScreenshot,
@@ -825,7 +827,7 @@
             try {
                 WindowManagerGlobal.getWindowManagerService()
                         .overridePendingAppTransitionRemote(runner,
-                                mDisplayTracker.getDefaultDisplayId());
+                                mDisplayId);
             } catch (Exception e) {
                 Log.e(TAG, "Error overriding screenshot app transition", e);
             }
@@ -1160,8 +1162,8 @@
         }
     }
 
-    private Display getDefaultDisplay() {
-        return mDisplayManager.getDisplay(mDisplayTracker.getDefaultDisplayId());
+    private Display getDisplay() {
+        return mDisplayManager.getDisplay(mDisplayId);
     }
 
     private boolean allowLongScreenshots() {
@@ -1170,7 +1172,7 @@
 
     private Rect getFullScreenRect() {
         DisplayMetrics displayMetrics = new DisplayMetrics();
-        getDefaultDisplay().getRealMetrics(displayMetrics);
+        getDisplay().getRealMetrics(displayMetrics);
         return new Rect(0, 0, displayMetrics.widthPixels, displayMetrics.heightPixels);
     }
 
@@ -1229,4 +1231,11 @@
             };
         }
     }
+
+    /** Injectable factory to create screenshot controller instances for a specific display. */
+    @AssistedFactory
+    public interface Factory {
+        /** Creates an instance of the controller for that specific displayId. */
+        ScreenshotController create(int displayId);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
index e9be88a..92e933a 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
@@ -6,12 +6,13 @@
 import android.graphics.Rect
 import android.net.Uri
 import android.os.UserHandle
+import android.view.Display
 import android.view.WindowManager.ScreenshotSource
 import android.view.WindowManager.ScreenshotType
 import androidx.annotation.VisibleForTesting
 import com.android.internal.util.ScreenshotRequest
 
-/** ScreenshotData represents the current state of a single screenshot being acquired. */
+/** [ScreenshotData] represents the current state of a single screenshot being acquired. */
 data class ScreenshotData(
     @ScreenshotType var type: Int,
     @ScreenshotSource var source: Int,
@@ -23,6 +24,7 @@
     var taskId: Int,
     var insets: Insets,
     var bitmap: Bitmap?,
+    var displayId: Int,
     /** App-provided URL representing the content the user was looking at in the screenshot. */
     var contextUrl: Uri? = null,
 ) {
@@ -31,22 +33,31 @@
 
     companion object {
         @JvmStatic
-        fun fromRequest(request: ScreenshotRequest): ScreenshotData {
-            return ScreenshotData(
-                request.type,
-                request.source,
-                if (request.userId >= 0) UserHandle.of(request.userId) else null,
-                request.topComponent,
-                request.boundsInScreen,
-                request.taskId,
-                request.insets,
-                request.bitmap,
+        fun fromRequest(request: ScreenshotRequest, displayId: Int = Display.DEFAULT_DISPLAY) =
+            ScreenshotData(
+                type = request.type,
+                source = request.source,
+                userHandle = if (request.userId >= 0) UserHandle.of(request.userId) else null,
+                topComponent = request.topComponent,
+                screenBounds = request.boundsInScreen,
+                taskId = request.taskId,
+                insets = request.insets,
+                bitmap = request.bitmap,
+                displayId = displayId,
             )
-        }
 
         @VisibleForTesting
-        fun forTesting(): ScreenshotData {
-            return ScreenshotData(0, 0, null, null, null, 0, Insets.NONE, null)
-        }
+        fun forTesting() =
+            ScreenshotData(
+                type = 0,
+                source = 0,
+                userHandle = null,
+                topComponent = null,
+                screenBounds = null,
+                taskId = 0,
+                insets = Insets.NONE,
+                bitmap = null,
+                displayId = Display.DEFAULT_DISPLAY,
+            )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
new file mode 100644
index 0000000..6c886fc
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
@@ -0,0 +1,201 @@
+package com.android.systemui.screenshot
+
+import android.net.Uri
+import android.os.Trace
+import android.util.Log
+import android.view.Display
+import com.android.internal.logging.UiEventLogger
+import com.android.internal.util.ScreenshotRequest
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.display.data.repository.DisplayRepository
+import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback
+import java.util.function.Consumer
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
+
+/**
+ * Receives the signal to take a screenshot from [TakeScreenshotService], and calls back with the
+ * result.
+ *
+ * Captures a screenshot for each [Display] available.
+ */
+@SysUISingleton
+class TakeScreenshotExecutor
+@Inject
+constructor(
+    private val screenshotControllerFactory: ScreenshotController.Factory,
+    displayRepository: DisplayRepository,
+    @Application private val mainScope: CoroutineScope,
+    private val screenshotRequestProcessor: ScreenshotRequestProcessor,
+    private val uiEventLogger: UiEventLogger
+) {
+
+    private lateinit var displays: StateFlow<Set<Display>>
+    private val displaysCollectionJob: Job =
+        mainScope.launch {
+            displays = displayRepository.displays.stateIn(this, SharingStarted.Eagerly, emptySet())
+        }
+
+    private val screenshotControllers = mutableMapOf<Int, ScreenshotController>()
+
+    /**
+     * Executes the [ScreenshotRequest].
+     *
+     * [onSaved] is invoked only on the default display result. [RequestCallback.onFinish] is
+     * invoked only when both screenshot UIs are removed.
+     */
+    suspend fun executeScreenshots(
+        screenshotRequest: ScreenshotRequest,
+        onSaved: (Uri) -> Unit,
+        requestCallback: RequestCallback
+    ) {
+        val displayIds = getDisplaysToScreenshot()
+        val resultCallbackWrapper = MultiResultCallbackWrapper(requestCallback)
+        screenshotRequest.oneForEachDisplay(displayIds).forEach { screenshotData: ScreenshotData ->
+            dispatchToController(
+                screenshotData = screenshotData,
+                onSaved =
+                    if (screenshotData.displayId == Display.DEFAULT_DISPLAY) onSaved else { _ -> },
+                callback = resultCallbackWrapper.createCallbackForId(screenshotData.displayId)
+            )
+        }
+    }
+
+    /** Creates a [ScreenshotData] for each display. */
+    private suspend fun ScreenshotRequest.oneForEachDisplay(
+        displayIds: List<Int>
+    ): List<ScreenshotData> {
+        return displayIds
+            .map { displayId -> ScreenshotData.fromRequest(this, displayId) }
+            .map { screenshotData: ScreenshotData ->
+                screenshotRequestProcessor.process(screenshotData)
+            }
+    }
+
+    private fun dispatchToController(
+        screenshotData: ScreenshotData,
+        onSaved: (Uri) -> Unit,
+        callback: RequestCallback
+    ) {
+        uiEventLogger.log(
+            ScreenshotEvent.getScreenshotSource(screenshotData.source),
+            0,
+            screenshotData.packageNameString
+        )
+        Log.d(TAG, "Screenshot request: $screenshotData")
+        getScreenshotController(screenshotData.displayId)
+            .handleScreenshot(screenshotData, onSaved, callback)
+    }
+
+    private fun getDisplaysToScreenshot(): List<Int> {
+        return displays.value.filter { it.type in ALLOWED_DISPLAY_TYPES }.map { it.displayId }
+    }
+
+    /**
+     * Propagates the close system dialog signal to all controllers.
+     *
+     * TODO(b/295143676): Move the receiver in this class once the flag is flipped.
+     */
+    fun onCloseSystemDialogsReceived() {
+        screenshotControllers.forEach { (_, screenshotController) ->
+            if (!screenshotController.isPendingSharedTransition) {
+                screenshotController.dismissScreenshot(ScreenshotEvent.SCREENSHOT_DISMISSED_OTHER)
+            }
+        }
+    }
+
+    /** Removes all screenshot related windows. */
+    fun removeWindows() {
+        screenshotControllers.forEach { (_, screenshotController) ->
+            screenshotController.removeWindow()
+        }
+    }
+
+    /**
+     * Destroys the executor. Afterwards, this class is not expected to work as intended anymore.
+     */
+    fun onDestroy() {
+        screenshotControllers.forEach { (_, screenshotController) ->
+            screenshotController.onDestroy()
+        }
+        screenshotControllers.clear()
+        displaysCollectionJob.cancel()
+    }
+
+    private fun getScreenshotController(id: Int): ScreenshotController {
+        return screenshotControllers.computeIfAbsent(id) { screenshotControllerFactory.create(id) }
+    }
+
+    /** For java compatibility only. see [executeScreenshots] */
+    fun executeScreenshotsAsync(
+        screenshotRequest: ScreenshotRequest,
+        onSaved: Consumer<Uri>,
+        requestCallback: RequestCallback
+    ) {
+        mainScope.launch {
+            executeScreenshots(screenshotRequest, { uri -> onSaved.accept(uri) }, requestCallback)
+        }
+    }
+
+    /**
+     * Returns a [RequestCallback] that calls [RequestCallback.onFinish] only when all callbacks for
+     * id created have finished.
+     *
+     * If any callback created calls [reportError], then following [onFinish] are not considered.
+     */
+    private class MultiResultCallbackWrapper(
+        private val originalCallback: RequestCallback,
+    ) {
+        private val idsPending = mutableSetOf<Int>()
+        private var errorReported = false
+
+        /**
+         * Creates a callback for [id].
+         *
+         * [originalCallback]'s [onFinish] will be called only when this (and the other created)
+         * callback's [onFinish] have been called.
+         */
+        fun createCallbackForId(id: Int): RequestCallback {
+            Trace.asyncTraceForTrackBegin(Trace.TRACE_TAG_APP, TAG, "Waiting for id=$id", id)
+            idsPending += id
+            return object : RequestCallback {
+                override fun reportError() {
+                    Log.d(TAG, "ReportError id=$id")
+                    Trace.asyncTraceForTrackEnd(Trace.TRACE_TAG_APP, TAG, id)
+                    Trace.instantForTrack(Trace.TRACE_TAG_APP, TAG, "reportError id=$id")
+                    originalCallback.reportError()
+                    errorReported = true
+                }
+
+                override fun onFinish() {
+                    Log.d(TAG, "onFinish id=$id")
+                    if (errorReported) return
+                    idsPending -= id
+                    Trace.instantForTrack(Trace.TRACE_TAG_APP, TAG, "onFinish id=$id")
+                    if (idsPending.isEmpty()) {
+                        Trace.asyncTraceForTrackEnd(Trace.TRACE_TAG_APP, TAG, id)
+                        originalCallback.onFinish()
+                    }
+                }
+            }
+        }
+    }
+
+    private companion object {
+        val TAG = LogConfig.logTag(TakeScreenshotService::class.java)
+
+        val ALLOWED_DISPLAY_TYPES =
+            listOf(
+                Display.TYPE_EXTERNAL,
+                Display.TYPE_INTERNAL,
+                Display.TYPE_OVERLAY,
+                Display.TYPE_WIFI
+            )
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
index 1cdad83..1e8542f 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
@@ -21,6 +21,7 @@
 
 import static com.android.internal.util.ScreenshotHelper.SCREENSHOT_MSG_PROCESS_COMPLETE;
 import static com.android.internal.util.ScreenshotHelper.SCREENSHOT_MSG_URI;
+import static com.android.systemui.flags.Flags.MULTI_DISPLAY_SCREENSHOT;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_CALLBACK;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_DISMISS;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_SERVICE;
@@ -46,6 +47,7 @@
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.util.Log;
+import android.view.Display;
 import android.widget.Toast;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -59,6 +61,7 @@
 import java.util.function.Consumer;
 
 import javax.inject.Inject;
+import javax.inject.Provider;
 
 public class TakeScreenshotService extends Service {
     private static final String TAG = logTag(TakeScreenshotService.class);
@@ -82,12 +85,17 @@
                 if (DEBUG_DISMISS) {
                     Log.d(TAG, "Received ACTION_CLOSE_SYSTEM_DIALOGS");
                 }
-                if (!mScreenshot.isPendingSharedTransition()) {
+                if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+                    // TODO(b/295143676): move receiver inside executor when the flag is enabled.
+                    mTakeScreenshotExecutor.get().onCloseSystemDialogsReceived();
+                } else if (!mScreenshot.isPendingSharedTransition()) {
                     mScreenshot.dismissScreenshot(SCREENSHOT_DISMISSED_OTHER);
                 }
             }
         }
     };
+    private final Provider<TakeScreenshotExecutor> mTakeScreenshotExecutor;
+
 
     /** Informs about coarse grained state of the Controller. */
     public interface RequestCallback {
@@ -99,16 +107,15 @@
     }
 
     @Inject
-    public TakeScreenshotService(ScreenshotController screenshotController, UserManager userManager,
-            DevicePolicyManager devicePolicyManager, UiEventLogger uiEventLogger,
-            ScreenshotNotificationsController notificationsController, Context context,
-            @Background Executor bgExecutor, FeatureFlags featureFlags,
-            RequestProcessor processor) {
+    public TakeScreenshotService(ScreenshotController.Factory screenshotControllerFactory,
+            UserManager userManager, DevicePolicyManager devicePolicyManager,
+            UiEventLogger uiEventLogger, ScreenshotNotificationsController notificationsController,
+            Context context, @Background Executor bgExecutor, FeatureFlags featureFlags,
+            RequestProcessor processor, Provider<TakeScreenshotExecutor> takeScreenshotExecutor) {
         if (DEBUG_SERVICE) {
             Log.d(TAG, "new " + this);
         }
         mHandler = new Handler(Looper.getMainLooper(), this::handleMessage);
-        mScreenshot = screenshotController;
         mUserManager = userManager;
         mDevicePolicyManager = devicePolicyManager;
         mUiEventLogger = uiEventLogger;
@@ -117,6 +124,12 @@
         mBgExecutor = bgExecutor;
         mFeatureFlags = featureFlags;
         mProcessor = processor;
+        mTakeScreenshotExecutor = takeScreenshotExecutor;
+        if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+            mScreenshot = null;
+        } else {
+            mScreenshot = screenshotControllerFactory.create(Display.DEFAULT_DISPLAY);
+        }
     }
 
     @Override
@@ -142,7 +155,11 @@
         if (DEBUG_SERVICE) {
             Log.d(TAG, "onUnbind");
         }
-        mScreenshot.removeWindow();
+        if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+            mTakeScreenshotExecutor.get().removeWindows();
+        } else {
+            mScreenshot.removeWindow();
+        }
         unregisterReceiver(mCloseSystemDialogs);
         return false;
     }
@@ -150,7 +167,11 @@
     @Override
     public void onDestroy() {
         super.onDestroy();
-        mScreenshot.onDestroy();
+        if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+            mTakeScreenshotExecutor.get().onDestroy();
+        } else {
+            mScreenshot.onDestroy();
+        }
         if (DEBUG_SERVICE) {
             Log.d(TAG, "onDestroy");
         }
@@ -218,10 +239,17 @@
         }
 
         Log.d(TAG, "Processing screenshot data");
-        ScreenshotData screenshotData = ScreenshotData.fromRequest(request);
+
+
+        ScreenshotData screenshotData = ScreenshotData.fromRequest(
+                request, Display.DEFAULT_DISPLAY);
         try {
-            mProcessor.processAsync(screenshotData,
-                    (data) -> dispatchToController(data, onSaved, callback));
+            if (mFeatureFlags.isEnabled(MULTI_DISPLAY_SCREENSHOT)) {
+                mTakeScreenshotExecutor.get().executeScreenshotsAsync(request, onSaved, callback);
+            } else {
+                mProcessor.processAsync(screenshotData, (data) ->
+                        dispatchToController(data, onSaved, callback));
+            }
         } catch (IllegalStateException e) {
             Log.e(TAG, "Failed to process screenshot request!", e);
             logFailedRequest(request);
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java b/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
index 22e238c0..7d17d4c 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
@@ -20,9 +20,11 @@
 
 import com.android.systemui.screenshot.ImageCapture;
 import com.android.systemui.screenshot.ImageCaptureImpl;
+import com.android.systemui.screenshot.RequestProcessor;
 import com.android.systemui.screenshot.ScreenshotPolicy;
 import com.android.systemui.screenshot.ScreenshotPolicyImpl;
 import com.android.systemui.screenshot.ScreenshotProxyService;
+import com.android.systemui.screenshot.ScreenshotRequestProcessor;
 import com.android.systemui.screenshot.TakeScreenshotService;
 import com.android.systemui.screenshot.appclips.AppClipsScreenshotHelperService;
 import com.android.systemui.screenshot.appclips.AppClipsService;
@@ -63,4 +65,8 @@
     @IntoMap
     @ClassKey(AppClipsService.class)
     abstract Service bindAppClipsService(AppClipsService service);
+
+    @Binds
+    abstract ScreenshotRequestProcessor bindScreenshotRequestProcessor(
+            RequestProcessor requestProcessor);
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt
index 43e9939..f8a8a68 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt
@@ -20,6 +20,7 @@
 import android.graphics.Insets
 import android.graphics.Rect
 import android.os.UserHandle
+import android.view.Display
 import android.view.WindowManager
 import com.android.internal.util.ScreenshotRequest
 import com.google.common.truth.Truth.assertThat
@@ -54,6 +55,16 @@
         assertThat(data.taskId).isEqualTo(taskId)
         assertThat(data.userHandle).isEqualTo(UserHandle.of(userId))
         assertThat(data.topComponent).isEqualTo(component)
+        assertThat(data.displayId).isEqualTo(Display.DEFAULT_DISPLAY)
+    }
+
+    @Test
+    fun testConstruction_notDefaultDisplayId() {
+        val request = ScreenshotRequest.Builder(type, source).build()
+
+        val data = ScreenshotData.fromRequest(request, displayId = 42)
+
+        assertThat(data.displayId).isEqualTo(42)
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
new file mode 100644
index 0000000..97c2ed4
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
@@ -0,0 +1,303 @@
+package com.android.systemui.screenshot
+
+import android.content.ComponentName
+import android.net.Uri
+import android.testing.AndroidTestingRunner
+import android.view.Display
+import android.view.Display.TYPE_EXTERNAL
+import android.view.Display.TYPE_INTERNAL
+import android.view.Display.TYPE_OVERLAY
+import android.view.Display.TYPE_VIRTUAL
+import android.view.Display.TYPE_WIFI
+import android.view.WindowManager
+import androidx.test.filters.SmallTest
+import com.android.internal.logging.testing.UiEventLoggerFake
+import com.android.internal.util.ScreenshotRequest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.display.data.repository.FakeDisplayRepository
+import com.android.systemui.display.data.repository.display
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.mockito.kotlinArgumentCaptor as ArgumentCaptor
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.nullable
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.never
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.verifyNoMoreInteractions
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+class TakeScreenshotExecutorTest : SysuiTestCase() {
+
+    private val controller0 = mock<ScreenshotController>()
+    private val controller1 = mock<ScreenshotController>()
+    private val controllerFactory = mock<ScreenshotController.Factory>()
+    private val callback = mock<TakeScreenshotService.RequestCallback>()
+
+    private val fakeDisplayRepository = FakeDisplayRepository()
+    private val requestProcessor = FakeRequestProcessor()
+    private val topComponent = ComponentName(mContext, TakeScreenshotExecutorTest::class.java)
+    private val testScope = TestScope(UnconfinedTestDispatcher())
+    private val eventLogger = UiEventLoggerFake()
+
+    private val screenshotExecutor =
+        TakeScreenshotExecutor(
+            controllerFactory,
+            fakeDisplayRepository,
+            testScope,
+            requestProcessor,
+            eventLogger,
+        )
+
+    @Before
+    fun setUp() {
+        whenever(controllerFactory.create(eq(0))).thenReturn(controller0)
+        whenever(controllerFactory.create(eq(1))).thenReturn(controller1)
+    }
+
+    @Test
+    fun executeScreenshots_severalDisplays_callsControllerForEachOne() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            verify(controllerFactory).create(eq(0))
+            verify(controllerFactory).create(eq(1))
+
+            val capturer = ArgumentCaptor<ScreenshotData>()
+
+            verify(controller0).handleScreenshot(capturer.capture(), any(), any())
+            assertThat(capturer.value.displayId).isEqualTo(0)
+            // OnSaved callback should be different.
+            verify(controller1).handleScreenshot(capturer.capture(), any(), any())
+            assertThat(capturer.value.displayId).isEqualTo(1)
+
+            assertThat(eventLogger.numLogs()).isEqualTo(2)
+            assertThat(eventLogger.get(0).eventId)
+                .isEqualTo(ScreenshotEvent.SCREENSHOT_REQUESTED_KEY_OTHER.id)
+            assertThat(eventLogger.get(0).packageName).isEqualTo(topComponent.packageName)
+            assertThat(eventLogger.get(1).eventId)
+                .isEqualTo(ScreenshotEvent.SCREENSHOT_REQUESTED_KEY_OTHER.id)
+            assertThat(eventLogger.get(1).packageName).isEqualTo(topComponent.packageName)
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_onlyVirtualDisplays_noInteractionsWithControllers() =
+        testScope.runTest {
+            setDisplays(display(TYPE_VIRTUAL, id = 0), display(TYPE_VIRTUAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            verifyNoMoreInteractions(controllerFactory)
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_allowedTypes_allCaptured() =
+        testScope.runTest {
+            whenever(controllerFactory.create(any())).thenReturn(controller0)
+
+            setDisplays(
+                display(TYPE_INTERNAL, id = 0),
+                display(TYPE_EXTERNAL, id = 1),
+                display(TYPE_OVERLAY, id = 2),
+                display(TYPE_WIFI, id = 3)
+            )
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            verify(controller0, times(4)).handleScreenshot(any(), any(), any())
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_reportsOnFinishedOnlyWhenBothFinished() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            val capturer0 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+            val capturer1 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+
+            verify(controller0).handleScreenshot(any(), any(), capturer0.capture())
+            verify(controller1).handleScreenshot(any(), any(), capturer1.capture())
+
+            verify(callback, never()).onFinish()
+
+            capturer0.value.onFinish()
+
+            verify(callback, never()).onFinish()
+
+            capturer1.value.onFinish()
+
+            verify(callback).onFinish()
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_doesNotReportFinishedIfOneFinishesOtherFails() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            val capturer0 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+            val capturer1 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+
+            verify(controller0).handleScreenshot(any(), any(), capturer0.capture())
+            verify(controller1).handleScreenshot(any(), nullable(), capturer1.capture())
+
+            verify(callback, never()).onFinish()
+
+            capturer0.value.onFinish()
+
+            verify(callback, never()).onFinish()
+
+            capturer1.value.reportError()
+
+            verify(callback, never()).onFinish()
+            verify(callback).reportError()
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_doesNotReportFinishedAfterOneFails() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            val capturer0 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+            val capturer1 = ArgumentCaptor<TakeScreenshotService.RequestCallback>()
+
+            verify(controller0).handleScreenshot(any(), any(), capturer0.capture())
+            verify(controller1).handleScreenshot(any(), any(), capturer1.capture())
+
+            verify(callback, never()).onFinish()
+
+            capturer0.value.reportError()
+
+            verify(callback, never()).onFinish()
+            verify(callback).reportError()
+
+            capturer1.value.onFinish()
+
+            verify(callback, never()).onFinish()
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun onDestroy_propagatedToControllers() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            screenshotExecutor.onDestroy()
+            verify(controller0).onDestroy()
+            verify(controller1).onDestroy()
+        }
+
+    @Test
+    fun removeWindows_propagatedToControllers() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            screenshotExecutor.removeWindows()
+            verify(controller0).removeWindow()
+            verify(controller1).removeWindow()
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun onCloseSystemDialogsReceived_propagatedToControllers() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            screenshotExecutor.onCloseSystemDialogsReceived()
+            verify(controller0).dismissScreenshot(any())
+            verify(controller1).dismissScreenshot(any())
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun onCloseSystemDialogsReceived_someControllerHavePendingTransitions() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            whenever(controller0.isPendingSharedTransition).thenReturn(true)
+            whenever(controller1.isPendingSharedTransition).thenReturn(false)
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
+
+            screenshotExecutor.onCloseSystemDialogsReceived()
+            verify(controller0, never()).dismissScreenshot(any())
+            verify(controller1).dismissScreenshot(any())
+
+            screenshotExecutor.onDestroy()
+        }
+
+    @Test
+    fun executeScreenshots_controllerCalledWithRequestProcessorReturnValue() =
+        testScope.runTest {
+            setDisplays(display(TYPE_INTERNAL, id = 0))
+            val screenshotRequest = createScreenshotRequest()
+            val toBeReturnedByProcessor = ScreenshotData.forTesting()
+            requestProcessor.toReturn = toBeReturnedByProcessor
+
+            val onSaved = { _: Uri -> }
+            screenshotExecutor.executeScreenshots(screenshotRequest, onSaved, callback)
+
+            assertThat(requestProcessor.processed)
+                .isEqualTo(ScreenshotData.fromRequest(screenshotRequest))
+
+            val capturer = ArgumentCaptor<ScreenshotData>()
+            verify(controller0).handleScreenshot(capturer.capture(), any(), any())
+            assertThat(capturer.value).isEqualTo(toBeReturnedByProcessor)
+
+            screenshotExecutor.onDestroy()
+        }
+
+    private suspend fun TestScope.setDisplays(vararg displays: Display) {
+        fakeDisplayRepository.emit(displays.toSet())
+        runCurrent()
+    }
+
+    private fun createScreenshotRequest() =
+        ScreenshotRequest.Builder(
+                WindowManager.TAKE_SCREENSHOT_FULLSCREEN,
+                WindowManager.ScreenshotSource.SCREENSHOT_KEY_OTHER
+            )
+            .setTopComponent(topComponent)
+            .build()
+
+    private class FakeRequestProcessor : ScreenshotRequestProcessor {
+        var processed: ScreenshotData? = null
+        var toReturn: ScreenshotData? = null
+
+        override suspend fun process(screenshot: ScreenshotData): ScreenshotData {
+            processed = screenshot
+            return toReturn ?: screenshot
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
index 77f7426..a08cda6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
@@ -27,6 +27,7 @@
 import android.os.UserHandle
 import android.os.UserManager
 import android.testing.AndroidTestingRunner
+import android.view.Display
 import android.view.WindowManager.ScreenshotSource.SCREENSHOT_KEY_OTHER
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import androidx.test.filters.SmallTest
@@ -34,6 +35,7 @@
 import com.android.internal.util.ScreenshotRequest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.flags.FakeFeatureFlags
+import com.android.systemui.flags.Flags.MULTI_DISPLAY_SCREENSHOT
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_CAPTURE_FAILED
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_REQUESTED_KEY_OTHER
 import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback
@@ -48,6 +50,7 @@
 import org.junit.runner.RunWith
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.isNull
+import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.doThrow
 import org.mockito.Mockito.times
@@ -60,6 +63,8 @@
 
     private val application = mock<Application>()
     private val controller = mock<ScreenshotController>()
+    private val controllerFactory = mock<ScreenshotController.Factory>()
+    private val takeScreenshotExecutor = mock<TakeScreenshotExecutor>()
     private val userManager = mock<UserManager>()
     private val requestProcessor = mock<RequestProcessor>()
     private val devicePolicyManager = mock<DevicePolicyManager>()
@@ -71,21 +76,11 @@
     private val flags = FakeFeatureFlags()
     private val topComponent = ComponentName(mContext, TakeScreenshotServiceTest::class.java)
 
-    private val service =
-        TakeScreenshotService(
-            controller,
-            userManager,
-            devicePolicyManager,
-            eventLogger,
-            notificationsController,
-            mContext,
-            Runnable::run,
-            flags,
-            requestProcessor
-        )
+    private lateinit var service: TakeScreenshotService
 
     @Before
     fun setUp() {
+        flags.set(MULTI_DISPLAY_SCREENSHOT, false)
         whenever(devicePolicyManager.resources).thenReturn(devicePolicyResourcesManager)
         whenever(
                 devicePolicyManager.getScreenCaptureDisabled(
@@ -95,6 +90,7 @@
             )
             .thenReturn(false)
         whenever(userManager.isUserUnlocked).thenReturn(true)
+        whenever(controllerFactory.create(any())).thenReturn(controller)
 
         // Stub request processor as a synchronous no-op for tests with the flag enabled
         doAnswer {
@@ -113,14 +109,7 @@
             .whenever(requestProcessor)
             .processAsync(/* screenshot= */ any(ScreenshotData::class.java), /* callback= */ any())
 
-        service.attach(
-            mContext,
-            /* thread = */ null,
-            /* className = */ null,
-            /* token = */ null,
-            application,
-            /* activityManager = */ null
-        )
+        service = createService()
     }
 
     @Test
@@ -146,7 +135,7 @@
 
         verify(controller, times(1))
             .handleScreenshot(
-                eq(ScreenshotData.fromRequest(request)),
+                eq(ScreenshotData.fromRequest(request, Display.DEFAULT_DISPLAY)),
                 /* onSavedListener = */ any(),
                 /* requestCallback = */ any()
             )
@@ -295,6 +284,74 @@
             failureEvent.packageName
         )
     }
+
+    @Test
+    fun takeScreenshotFullScreen_multiDisplayFlagEnabled_takeScreenshotExecutor() {
+        flags.set(MULTI_DISPLAY_SCREENSHOT, true)
+        service = createService()
+
+        val request =
+            ScreenshotRequest.Builder(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_KEY_OTHER)
+                .setTopComponent(topComponent)
+                .build()
+
+        service.handleRequest(request, { /* onSaved */}, callback)
+
+        verifyZeroInteractions(controller)
+        verify(takeScreenshotExecutor, times(1)).executeScreenshotsAsync(any(), any(), any())
+
+        assertEquals("Expected one UiEvent", 0, eventLogger.numLogs())
+    }
+
+    @Test
+    fun testServiceLifecycle_multiDisplayScreenshotFlagEnabled() {
+        flags.set(MULTI_DISPLAY_SCREENSHOT, true)
+        service = createService()
+
+        service.onCreate()
+        service.onBind(null /* unused: Intent */)
+
+        service.onUnbind(null /* unused: Intent */)
+        verify(takeScreenshotExecutor, times(1)).removeWindows()
+
+        service.onDestroy()
+        verify(takeScreenshotExecutor, times(1)).onDestroy()
+    }
+
+    @Test
+    fun constructor_MultiDisplayFlagOn_screenshotControllerNotCreated() {
+        flags.set(MULTI_DISPLAY_SCREENSHOT, true)
+        clearInvocations(controllerFactory)
+
+        service = createService()
+
+        verifyZeroInteractions(controllerFactory)
+    }
+
+    private fun createService(): TakeScreenshotService {
+        val service =
+            TakeScreenshotService(
+                controllerFactory,
+                userManager,
+                devicePolicyManager,
+                eventLogger,
+                notificationsController,
+                mContext,
+                Runnable::run,
+                flags,
+                requestProcessor,
+                { takeScreenshotExecutor },
+            )
+        service.attach(
+            mContext,
+            /* thread = */ null,
+            /* className = */ null,
+            /* token = */ null,
+            application,
+            /* activityManager = */ null
+        )
+        return service
+    }
 }
 
 private fun Bitmap.equalsHardwareBitmap(other: Bitmap): Boolean {
diff --git a/services/core/java/com/android/server/pm/PackageManagerService.java b/services/core/java/com/android/server/pm/PackageManagerService.java
index f0bbd35..7ccf713 100644
--- a/services/core/java/com/android/server/pm/PackageManagerService.java
+++ b/services/core/java/com/android/server/pm/PackageManagerService.java
@@ -2892,10 +2892,13 @@
 
     private void notifyPackageUseInternal(String packageName, int reason) {
         long time = System.currentTimeMillis();
-        this.commitPackageStateMutation(null, mutator -> {
-            final PackageStateWrite state = mutator.forPackage(packageName);
-            state.setLastPackageUsageTime(reason, time);
-        });
+        synchronized (mLock) {
+            final PackageSetting pkgSetting = mSettings.getPackageLPr(packageName);
+            if (pkgSetting == null) {
+                return;
+            }
+            pkgSetting.getPkgState().setLastPackageUsageTimeInMills(reason, time);
+        }
     }
 
     /*package*/ DexManager getDexManager() {