Merge "Improved screenshot policy for desktop and split mode" into main
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
index 254f1e1..4d71dc4 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
@@ -21,8 +21,8 @@
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FREE_FORM
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.PIP
-import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.SPLIT_BOTTOM
-import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.SPLIT_TOP
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Orientation.HORIZONTAL
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Orientation.VERTICAL
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks.emptyRootSplit
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks.freeForm
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks.fullScreen
@@ -39,16 +39,14 @@
 
     data class TaskSpec(val taskId: Int, val userId: Int, val name: String)
 
+    val emptyDisplayContent = DisplayContentModel(0, SystemUiState(shadeExpanded = false), listOf())
+
     /** Home screen, with only the launcher visible */
     fun launcherOnly(shadeExpanded: Boolean = false) =
         DisplayContentModel(
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
-            rootTasks =
-                listOf(
-                    launcher(visible = true),
-                    emptyRootSplit,
-                )
+            rootTasks = listOf(launcher(visible = true), emptyRootSplit),
         )
 
     /** A Full screen activity for the personal (primary) user, with launcher behind it */
@@ -57,48 +55,72 @@
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
             rootTasks =
-                listOf(
-                    fullScreen(spec, visible = true),
-                    launcher(visible = false),
-                    emptyRootSplit,
-                )
+                listOf(fullScreen(spec, visible = true), launcher(visible = false), emptyRootSplit),
         )
 
+    enum class Orientation {
+        HORIZONTAL,
+        VERTICAL,
+    }
+
+    internal fun Rect.splitLeft(margin: Int = 0) = Rect(left, top, centerX() - margin, bottom)
+
+    internal fun Rect.splitRight(margin: Int = 0) = Rect(centerX() + margin, top, right, bottom)
+
+    internal fun Rect.splitTop(margin: Int = 0) = Rect(left, top, right, centerY() - margin)
+
+    internal fun Rect.splitBottom(margin: Int = 0) = Rect(left, centerY() + margin, right, bottom)
+
     fun splitScreenApps(
-        top: TaskSpec,
-        bottom: TaskSpec,
+        displayId: Int = 0,
+        parentBounds: Rect = FULL_SCREEN,
+        taskMargin: Int = 0,
+        orientation: Orientation = VERTICAL,
+        first: TaskSpec,
+        second: TaskSpec,
         focusedTaskId: Int,
+        parentTaskId: Int = 2,
         shadeExpanded: Boolean = false,
     ): DisplayContentModel {
-        val topBounds = SPLIT_TOP
-        val bottomBounds = SPLIT_BOTTOM
+
+        val firstBounds =
+            when (orientation) {
+                VERTICAL -> parentBounds.splitTop(taskMargin)
+                HORIZONTAL -> parentBounds.splitLeft(taskMargin)
+            }
+        val secondBounds =
+            when (orientation) {
+                VERTICAL -> parentBounds.splitBottom(taskMargin)
+                HORIZONTAL -> parentBounds.splitRight(taskMargin)
+            }
+
         return DisplayContentModel(
-            displayId = 0,
+            displayId = displayId,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
             rootTasks =
                 listOf(
                     newRootTaskInfo(
-                        taskId = 2,
+                        taskId = parentTaskId,
                         userId = TestUserIds.PERSONAL,
-                        bounds = FULL_SCREEN,
+                        bounds = parentBounds,
                         topActivity =
                             ComponentName.unflattenFromString(
-                                if (top.taskId == focusedTaskId) top.name else bottom.name
+                                if (first.taskId == focusedTaskId) first.name else second.name
                             ),
                     ) {
                         listOf(
                                 newChildTask(
-                                    taskId = top.taskId,
-                                    bounds = topBounds,
-                                    userId = top.userId,
-                                    name = top.name
+                                    taskId = first.taskId,
+                                    bounds = firstBounds,
+                                    userId = first.userId,
+                                    name = first.name,
                                 ),
                                 newChildTask(
-                                    taskId = bottom.taskId,
-                                    bounds = bottomBounds,
-                                    userId = bottom.userId,
-                                    name = bottom.name
-                                )
+                                    taskId = second.taskId,
+                                    bounds = secondBounds,
+                                    userId = second.userId,
+                                    name = second.name,
+                                ),
                             )
                             // Child tasks are ordered bottom-up in RootTaskInfo.
                             // Sort 'focusedTaskId' last.
@@ -106,7 +128,7 @@
                             .sortedBy { it.id == focusedTaskId }
                     },
                     launcher(visible = false),
-                )
+                ),
         )
     }
 
@@ -124,7 +146,7 @@
                     fullScreen?.also { add(fullScreen(it, visible = true)) }
                     add(launcher(visible = (fullScreen == null)))
                     add(emptyRootSplit)
-                }
+                },
         )
     }
 
@@ -142,7 +164,7 @@
         return DisplayContentModel(
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
-            rootTasks = freeFormTasks + launcher(visible = true) + emptyRootSplit
+            rootTasks = freeFormTasks + launcher(visible = true) + emptyRootSplit,
         )
     }
 
@@ -153,11 +175,18 @@
      * somewhat sensible in terms of logical position (Re: PIP, SPLIT, etc).
      */
     object Bounds {
+        // "Phone" size
         val FULL_SCREEN = Rect(0, 0, 1080, 2400)
         val PIP = Rect(440, 1458, 1038, 1794)
         val SPLIT_TOP = Rect(0, 0, 1080, 1187)
         val SPLIT_BOTTOM = Rect(0, 1213, 1080, 2400)
         val FREE_FORM = Rect(119, 332, 1000, 1367)
+
+        // "Tablet" size
+        val FREEFORM_FULL_SCREEN = Rect(0, 0, 2560, 1600)
+        val FREEFORM_MAXIMIZED = Rect(0, 48, 2560, 1480)
+        val FREEFORM_SPLIT_LEFT = Rect(0, 0, 1270, 1600)
+        val FREEFORM_SPLIT_RIGHT = Rect(1290, 0, 2560, 1600)
     }
 
     /** A collection of task names used in test scenarios */
@@ -177,6 +206,8 @@
             "com.google.android.youtube/" +
                 "com.google.android.apps.youtube.app.watchwhile.WatchWhileActivity"
 
+        const val MESSAGES = "com.google.android.apps.messaging/.ui.ConversationListActivity"
+
         /** The NexusLauncher activity */
         const val LAUNCHER =
             "com.google.android.apps.nexuslauncher/" +
@@ -220,7 +251,7 @@
             }
 
         /** NexusLauncher on the default display. Usually below all other visible tasks */
-        fun launcher(visible: Boolean) =
+        fun launcher(visible: Boolean, bounds: Rect = FULL_SCREEN) =
             newRootTaskInfo(
                 taskId = 1,
                 activityType = ActivityType.Home,
@@ -229,43 +260,63 @@
                 topActivity = ComponentName.unflattenFromString(ActivityNames.LAUNCHER),
                 topActivityType = ActivityType.Home,
             ) {
-                listOf(newChildTask(taskId = 1002, name = ActivityNames.LAUNCHER))
+                listOf(newChildTask(taskId = 1002, name = ActivityNames.LAUNCHER, bounds = bounds))
             }
 
         /** A full screen Activity */
-        fun fullScreen(task: TaskSpec, visible: Boolean) =
+        fun fullScreen(task: TaskSpec, visible: Boolean, bounds: Rect = FULL_SCREEN) =
             newRootTaskInfo(
                 taskId = task.taskId,
                 userId = task.userId,
                 visible = visible,
-                bounds = FULL_SCREEN,
+                bounds = bounds,
                 topActivity = ComponentName.unflattenFromString(task.name),
             ) {
-                listOf(newChildTask(taskId = task.taskId, userId = task.userId, name = task.name))
+                listOf(
+                    newChildTask(
+                        taskId = task.taskId,
+                        userId = task.userId,
+                        name = task.name,
+                        bounds = bounds,
+                    )
+                )
             }
 
         /** An activity in Picture-in-Picture mode */
-        fun pictureInPicture(task: TaskSpec) =
+        fun pictureInPicture(task: TaskSpec, bounds: Rect = PIP) =
             newRootTaskInfo(
                 taskId = task.taskId,
                 userId = task.userId,
-                bounds = PIP,
                 windowingMode = WindowingMode.PictureInPicture,
                 topActivity = ComponentName.unflattenFromString(task.name),
             ) {
-                listOf(newChildTask(taskId = task.taskId, userId = userId, name = task.name))
+                listOf(
+                    newChildTask(
+                        taskId = task.taskId,
+                        userId = userId,
+                        name = task.name,
+                        bounds = bounds,
+                    )
+                )
             }
 
         /** An activity in FreeForm mode */
-        fun freeForm(task: TaskSpec) =
+        fun freeForm(task: TaskSpec, bounds: Rect = FREE_FORM) =
             newRootTaskInfo(
                 taskId = task.taskId,
                 userId = task.userId,
-                bounds = FREE_FORM,
+                bounds = bounds,
                 windowingMode = WindowingMode.Freeform,
                 topActivity = ComponentName.unflattenFromString(task.name),
             ) {
-                listOf(newChildTask(taskId = task.taskId, userId = userId, name = task.name))
+                listOf(
+                    newChildTask(
+                        taskId = task.taskId,
+                        userId = userId,
+                        name = task.name,
+                        bounds = bounds,
+                    )
+                )
             }
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
index 6c35b23..cedf0c8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
@@ -69,7 +69,7 @@
     taskId: Int,
     name: String,
     bounds: Rect? = null,
-    userId: Int? = null
+    userId: Int? = null,
 ): ChildTaskModel {
     return ChildTaskModel(taskId, name, bounds ?: this.bounds, userId ?: this.userId)
 }
@@ -83,7 +83,7 @@
     running: Boolean = true,
     activityType: ActivityType = Standard,
     windowingMode: WindowingMode = FullScreen,
-    bounds: Rect? = null,
+    bounds: Rect = Rect(),
     topActivity: ComponentName? = null,
     topActivityType: ActivityType = Standard,
     numActivities: Int? = null,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
index 6e57761..b7f565d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
@@ -17,8 +17,8 @@
 package com.android.systemui.screenshot.policy
 
 import android.content.ComponentName
-import androidx.test.ext.junit.runners.AndroidJUnit4
 import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.screenshot.data.model.DisplayContentModel
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
@@ -59,7 +59,7 @@
             policy.check(
                 singleFullScreen(
                     spec = TaskSpec(taskId = 1002, name = YOUTUBE, userId = PRIVATE),
-                    shadeExpanded = true
+                    shadeExpanded = true,
                 )
             )
 
@@ -93,8 +93,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -110,25 +110,20 @@
                         listOf(
                             fullScreen(
                                 TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
-                                visible = true
+                                visible = true,
                             ),
                             fullScreen(
                                 TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
-                                visible = false
+                                visible = false,
                             ),
                             launcher(visible = false),
                             emptyRootSplit,
-                        )
+                        ),
                 )
             )
 
         assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    PrivateProfilePolicy.NAME,
-                    PrivateProfilePolicy.NO_VISIBLE_TASKS,
-                )
-            )
+            .isEqualTo(NotMatched(PrivateProfilePolicy.NAME, PrivateProfilePolicy.NO_VISIBLE_TASKS))
     }
 
     @Test
@@ -136,9 +131,9 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
-                    focusedTaskId = 1003
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1003,
                 )
             )
 
@@ -150,8 +145,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -161,9 +156,9 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
-                    focusedTaskId = 1002
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1002,
                 )
             )
 
@@ -175,8 +170,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(FILES),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -196,8 +191,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE_PIP),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -220,8 +215,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE_PIP),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/ScreenshotPolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/ScreenshotPolicyTest.kt
new file mode 100644
index 0000000..28eb9fc
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/ScreenshotPolicyTest.kt
@@ -0,0 +1,351 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.screenshot.policy
+
+import android.content.ComponentName
+import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.LAUNCHER
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.MESSAGES
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.YOUTUBE
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FREEFORM_FULL_SCREEN
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Orientation.VERTICAL
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.freeFormApps
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.pictureInPictureApp
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.splitScreenApps
+import com.android.systemui.screenshot.data.repository.profileTypeRepository
+import com.android.systemui.screenshot.policy.CaptureType.FullScreen
+import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
+import com.android.systemui.screenshot.policy.CaptureType.RootTask
+import com.android.systemui.screenshot.policy.TestUserIds.PERSONAL
+import com.android.systemui.screenshot.policy.TestUserIds.PRIVATE
+import com.android.systemui.screenshot.policy.TestUserIds.WORK
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class ScreenshotPolicyTest {
+    private val kosmos = Kosmos()
+
+    private val defaultComponent = ComponentName("default", "default")
+    private val defaultOwner = UserHandle.SYSTEM
+
+    @Test
+    fun fullScreen_work() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                singleFullScreen(TaskSpec(taskId = 1002, name = FILES, userId = WORK)),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN),
+                    component = ComponentName.unflattenFromString(FILES),
+                    owner = UserHandle.of(WORK),
+                )
+            )
+    }
+
+    @Test
+    fun fullScreen_private() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                singleFullScreen(TaskSpec(taskId = 1002, name = YOUTUBE, userId = PRIVATE)),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_workAndPersonal() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PERSONAL),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_personalAndPrivate() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_workAndPrivate() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_twoWorkTasks() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    parentTaskId = 1,
+                    parentBounds = FREEFORM_FULL_SCREEN,
+                    orientation = VERTICAL,
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = WORK),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type =
+                        RootTask(
+                            parentTaskId = 1,
+                            taskBounds = FREEFORM_FULL_SCREEN,
+                            childTaskIds = listOf(1002, 1003),
+                        ),
+                    component = ComponentName.unflattenFromString(FILES),
+                    owner = UserHandle.of(WORK),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floatingWindows() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1003,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PERSONAL),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floatingWindows_maximized() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1003,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PERSONAL),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floatingWindows_withPrivate() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    TaskSpec(taskId = 1004, name = MESSAGES, userId = PERSONAL),
+                    focusedTaskId = 1004,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floating_workOnly() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(LAUNCHER),
+                    owner = defaultOwner,
+                )
+            )
+    }
+
+    @Test
+    fun fullScreen_shadeExpanded() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                singleFullScreen(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    shadeExpanded = true,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = defaultComponent,
+                    owner = defaultOwner,
+                )
+            )
+    }
+
+    @Test
+    fun fullScreen_with_PictureInPicture() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                pictureInPictureApp(
+                    pip = TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
+                    fullScreen = TaskSpec(taskId = 1003, name = FILES, userId = WORK),
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = IsolatedTask(taskId = 1003, taskBounds = FULL_SCREEN),
+                    component = ComponentName.unflattenFromString(FILES),
+                    owner = UserHandle.of(WORK),
+                )
+            )
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
index be9fcc2..30a786c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
@@ -31,13 +31,13 @@
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.YOUTUBE
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FREE_FORM
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
-import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.SPLIT_TOP
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.freeFormApps
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.pictureInPictureApp
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.splitScreenApps
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.splitTop
 import com.android.systemui.screenshot.data.model.SystemUiState
 import com.android.systemui.screenshot.data.repository.profileTypeRepository
 import com.android.systemui.screenshot.policy.CapturePolicy.PolicyResult
@@ -69,6 +69,7 @@
     @JvmField @Rule(order = 2) val mockitoRule: MockitoRule = MockitoJUnit.rule()
 
     @Mock lateinit var mContext: Context
+
     @Mock lateinit var mResources: Resources
 
     private val kosmos = Kosmos()
@@ -94,17 +95,11 @@
                 DisplayContentModel(
                     displayId = 0,
                     systemUiState = SystemUiState(shadeExpanded = false),
-                    rootTasks = listOf(RootTasks.emptyWithNoChildTasks)
+                    rootTasks = listOf(RootTasks.emptyWithNoChildTasks),
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    WORK_TASK_NOT_TOP,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, WORK_TASK_NOT_TOP))
     }
 
     @Test
@@ -114,13 +109,7 @@
                 singleFullScreen(TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL))
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    WORK_TASK_NOT_TOP,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, WORK_TASK_NOT_TOP))
     }
 
     @Test
@@ -129,17 +118,11 @@
             policy.check(
                 singleFullScreen(
                     TaskSpec(taskId = 1002, name = FILES, userId = WORK),
-                    shadeExpanded = true
+                    shadeExpanded = true,
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    SHADE_EXPANDED,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, SHADE_EXPANDED))
     }
 
     @Test
@@ -156,7 +139,7 @@
                         type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -166,9 +149,11 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
-                    focusedTaskId = 1002
+                    parentBounds = FULL_SCREEN,
+                    taskMargin = 20,
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1002,
                 )
             )
 
@@ -178,10 +163,10 @@
                     policy = WorkProfilePolicy.NAME,
                     reason = WORK_TASK_IS_TOP,
                     CaptureParameters(
-                        type = IsolatedTask(taskId = 1002, taskBounds = SPLIT_TOP),
+                        type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN.splitTop(20)),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -191,19 +176,13 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
-                    focusedTaskId = 1003
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1003,
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    WORK_TASK_NOT_TOP,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, WORK_TASK_NOT_TOP))
     }
 
     @Test
@@ -225,7 +204,7 @@
                         type = IsolatedTask(taskId = 1003, taskBounds = FULL_SCREEN),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -238,7 +217,7 @@
                 freeFormApps(
                     TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
                     TaskSpec(taskId = 1003, name = FILES, userId = WORK),
-                    focusedTaskId = 1003
+                    focusedTaskId = 1003,
                 )
             )
 
@@ -251,7 +230,7 @@
                         type = IsolatedTask(taskId = 1003, taskBounds = FREE_FORM),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -264,16 +243,10 @@
                 freeFormApps(
                     TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
                     TaskSpec(taskId = 1003, name = FILES, userId = WORK),
-                    focusedTaskId = 1003
+                    focusedTaskId = 1003,
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    DESKTOP_MODE_ENABLED,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, DESKTOP_MODE_ENABLED))
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt
index 0ef5207..9455201 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt
@@ -24,8 +24,8 @@
     data class FullScreen(val displayId: Int) : CaptureType
 
     /** Capture the contents of the task only. */
-    data class IsolatedTask(
-        val taskId: Int,
-        val taskBounds: Rect?,
-    ) : CaptureType
+    data class IsolatedTask(val taskId: Int, val taskBounds: Rect?) : CaptureType
+
+    data class RootTask(val parentTaskId: Int, val taskBounds: Rect?, val childTaskIds: List<Int>) :
+        CaptureType
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
index 039143a..e840668 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
@@ -26,6 +26,7 @@
 import android.util.Log
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
+import com.android.systemui.Flags.screenshotPolicySplitAndDesktopMode
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.screenshot.ImageCapture
 import com.android.systemui.screenshot.ScreenshotData
@@ -47,14 +48,17 @@
     private val capture: ImageCapture,
     /** Provides information about the tasks on a given display */
     private val displayTasks: DisplayContentRepository,
-    /** The list of policies to apply, in order of priority */
+    /** The legacy list of policy implementations to apply, in order of priority */
     private val policies: List<CapturePolicy>,
+    /** Implements the combined policy rules for all profile types. */
+    private val policy: ScreenshotPolicy,
     /** The owner to assign for screenshot when a focused task isn't visible */
     private val defaultOwner: UserHandle = myUserHandle(),
     /** The assigned component when no application has focus, or not visible */
     private val defaultComponent: ComponentName,
 ) : ScreenshotRequestProcessor {
     override suspend fun process(original: ScreenshotData): ScreenshotData {
+
         if (original.type == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
             // The request contains an already captured screenshot, accept it as is.
             Log.i(TAG, "Screenshot bitmap provided. No modifications applied.")
@@ -62,6 +66,12 @@
         }
         val displayContent = displayTasks.getDisplayContent(original.displayId)
 
+        if (screenshotPolicySplitAndDesktopMode()) {
+            Log.i(TAG, "Applying screenshot policy....")
+            val type = policy.apply(displayContent, defaultComponent, defaultOwner)
+            return modify(original, type)
+        }
+
         // If policies yield explicit modifications, apply them and return the result
         Log.i(TAG, "Applying policy checks....")
         policies.map { policy ->
@@ -79,10 +89,8 @@
     }
 
     /** Produce a new [ScreenshotData] using [CaptureParameters] */
-    private suspend fun modify(
-        original: ScreenshotData,
-        updates: CaptureParameters,
-    ): ScreenshotData {
+    suspend fun modify(original: ScreenshotData, updates: CaptureParameters): ScreenshotData {
+        Log.d(TAG, "[modify] CaptureParameters = $updates")
         // Update and apply bitmap capture depending on the parameters.
         val updated =
             when (val type = updates.type) {
@@ -94,6 +102,14 @@
                         type.taskId,
                         type.taskBounds,
                     )
+                is CaptureType.RootTask ->
+                    replaceWithTaskSnapshot(
+                        original,
+                        updates.component,
+                        updates.owner,
+                        type.parentTaskId,
+                        type.taskBounds,
+                    )
                 is FullScreen ->
                     replaceWithScreenshot(
                         original,
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt
index f768cfb..dd39f92 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt
@@ -26,9 +26,11 @@
             childTaskIds[index],
             childTaskNames[index],
             childTaskBounds[index],
-            childTaskUserIds[index]
+            childTaskUserIds[index],
         )
     }
 }
 
 internal fun RootTaskInfo.hasChildTasks() = childTaskUserIds.isNotEmpty()
+
+internal fun RootTaskInfo.childTaskCount() = childTaskIds.size
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicy.kt
new file mode 100644
index 0000000..9967aff
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicy.kt
@@ -0,0 +1,155 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.screenshot.policy
+
+import android.app.ActivityTaskManager.RootTaskInfo
+import android.app.WindowConfiguration
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WINDOWING_MODE_PINNED
+import android.content.ComponentName
+import android.os.UserHandle
+import android.util.Log
+import com.android.systemui.screenshot.data.model.DisplayContentModel
+import com.android.systemui.screenshot.data.model.ProfileType
+import com.android.systemui.screenshot.data.model.ProfileType.PRIVATE
+import com.android.systemui.screenshot.data.model.ProfileType.WORK
+import com.android.systemui.screenshot.data.repository.ProfileTypeRepository
+import com.android.systemui.screenshot.policy.CaptureType.FullScreen
+import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
+import com.android.systemui.screenshot.policy.CaptureType.RootTask
+import javax.inject.Inject
+
+private const val TAG = "ScreenshotPolicy"
+
+/** Determines what to capture and which user owns the output. */
+class ScreenshotPolicy @Inject constructor(private val profileTypes: ProfileTypeRepository) {
+    /**
+     * Apply the policy to the content, resulting in [CaptureParameters].
+     *
+     * @param content the content of the display
+     * @param defaultComponent the component associated with the screenshot by default
+     * @param defaultOwner the user to own the screenshot by default
+     */
+    suspend fun apply(
+        content: DisplayContentModel,
+        defaultComponent: ComponentName,
+        defaultOwner: UserHandle,
+    ): CaptureParameters {
+        val defaultFullScreen by lazy {
+            CaptureParameters(
+                type = FullScreen(displayId = content.displayId),
+                component = defaultComponent,
+                owner = defaultOwner,
+            )
+        }
+
+        // When the systemUI notification shade is open, disregard tasks.
+        if (content.systemUiState.shadeExpanded) {
+            return defaultFullScreen
+        }
+
+        // find the first (top) RootTask which is visible and not Picture-in-Picture
+        val topRootTask =
+            content.rootTasks.firstOrNull {
+                it.isVisible && it.windowingMode != WindowConfiguration.WINDOWING_MODE_PINNED
+            } ?: return defaultFullScreen
+
+        Log.d(TAG, "topRootTask: $topRootTask")
+        val rootTaskOwners = topRootTask.childTaskUserIds.distinct()
+
+        // Special case: Only WORK in top root task which is full-screen or maximized freeform
+        if (
+            rootTaskOwners.size == 1 &&
+                profileTypes.getProfileType(rootTaskOwners.single()) == WORK &&
+                (topRootTask.isFullScreen() || topRootTask.isMaximizedFreeform())
+        ) {
+            val type =
+                if (topRootTask.childTaskCount() > 1) {
+                    RootTask(
+                        parentTaskId = topRootTask.taskId,
+                        taskBounds = topRootTask.bounds,
+                        childTaskIds = topRootTask.childTasksTopDown().map { it.id }.toList(),
+                    )
+                } else {
+                    IsolatedTask(
+                        taskId = topRootTask.childTasksTopDown().first().id,
+                        taskBounds = topRootTask.bounds,
+                    )
+                }
+            // Capture the RootTask (and all children)
+            return CaptureParameters(
+                type = type,
+                component = topRootTask.topActivity,
+                owner = UserHandle.of(rootTaskOwners.single()),
+            )
+        }
+
+        // In every other case the output will be a full screen capture regardless of content.
+        // For this reason, consider all owners of all visible content on the display (in all
+        // root tasks). This includes all root tasks in free-form mode.
+        val visibleChildTasks =
+            content.rootTasks.filter { it.isVisible }.flatMap { it.childTasksTopDown() }
+
+        val allVisibleProfileTypes =
+            visibleChildTasks
+                .map { it.userId }
+                .distinct()
+                .associate { profileTypes.getProfileType(it) to UserHandle.of(it) }
+
+        // If any visible content belongs to the private profile user -> private profile
+        // otherwise the personal user (including partial screen work content).
+        val ownerHandle =
+            allVisibleProfileTypes[PRIVATE]
+                ?: allVisibleProfileTypes[ProfileType.NONE]
+                ?: defaultOwner
+
+        // Attribute to the component of top-most task owned by this user (or fallback to default)
+        val topComponent =
+            visibleChildTasks.firstOrNull { it.userId == ownerHandle.identifier }?.componentName
+
+        return CaptureParameters(
+            type = FullScreen(content.displayId),
+            component = topComponent ?: topRootTask.topActivity ?: defaultComponent,
+            owner = ownerHandle,
+        )
+    }
+
+    private fun RootTaskInfo.isFullScreen(): Boolean =
+        configuration.windowConfiguration.windowingMode == WINDOWING_MODE_FULLSCREEN
+
+    private fun RootTaskInfo.isMaximizedFreeform(): Boolean {
+        val bounds = configuration.windowConfiguration.bounds
+        val maxBounds = configuration.windowConfiguration.maxBounds
+
+        if (
+            windowingMode != WINDOWING_MODE_FREEFORM ||
+                childTaskCount() != 1 ||
+                childTaskBounds[0] != bounds
+        ) {
+            return false
+        }
+
+        // Maximized floating windows fill maxBounds width
+        if (bounds.width() != maxBounds.width()) {
+            return false
+        }
+
+        // Maximized floating windows fill nearly all the height
+        return (bounds.height().toFloat() / maxBounds.height()) >= 0.89f
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
index 2cb9fe7..a9c6370 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
@@ -37,7 +37,6 @@
 
 @Module
 interface ScreenshotPolicyModule {
-
     @Binds
     @SysUISingleton
     fun bindProfileTypeRepository(impl: ProfileTypeRepositoryImpl): ProfileTypeRepository
@@ -67,6 +66,7 @@
             imageCapture: ImageCapture,
             displayContentRepo: DisplayContentRepository,
             policyListProvider: Provider<List<CapturePolicy>>,
+            standardPolicy: ScreenshotPolicy,
         ): ScreenshotRequestProcessor {
             return PolicyRequestProcessor(
                 background = background,
@@ -75,7 +75,8 @@
                 policies = policyListProvider.get(),
                 defaultOwner = Process.myUserHandle(),
                 defaultComponent =
-                    ComponentName(context.packageName, SystemUIService::class.java.toString())
+                    ComponentName(context.packageName, SystemUIService::class.java.toString()),
+                policy = standardPolicy,
             )
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
index 29450a2..cf90c0a 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
@@ -28,7 +28,6 @@
 import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import javax.inject.Inject
-import kotlinx.coroutines.flow.first
 
 /**
  * Condition: When the top visible task (excluding PIP mode) belongs to a work user.
@@ -37,10 +36,8 @@
  */
 class WorkProfilePolicy
 @Inject
-constructor(
-    private val profileTypes: ProfileTypeRepository,
-    private val context: Context,
-) : CapturePolicy {
+constructor(private val profileTypes: ProfileTypeRepository, private val context: Context) :
+    CapturePolicy {
 
     override suspend fun check(content: DisplayContentModel): PolicyResult {
         // The systemUI notification shade isn't a work app, skip.
@@ -65,11 +62,7 @@
                 .map { it to it.childTasksTopDown().first() }
                 .firstOrNull { (_, child) ->
                     profileTypes.getProfileType(child.userId) == ProfileType.WORK
-                }
-                ?: return NotMatched(
-                    policy = NAME,
-                    reason = WORK_TASK_NOT_TOP,
-                )
+                } ?: return NotMatched(policy = NAME, reason = WORK_TASK_NOT_TOP)
 
         // If matched, return parameters needed to modify the request.
         return PolicyResult.Matched(
@@ -79,7 +72,7 @@
                 type = IsolatedTask(taskId = childTask.id, taskBounds = childTask.bounds),
                 component = childTask.componentName ?: rootTask.topActivity,
                 owner = UserHandle.of(childTask.userId),
-            )
+            ),
         )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
index 0d4cb4c..7709a65 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
@@ -20,30 +20,141 @@
 import android.graphics.Insets
 import android.graphics.Rect
 import android.os.UserHandle
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import android.view.Display.DEFAULT_DISPLAY
 import android.view.WindowManager.ScreenshotSource.SCREENSHOT_KEY_CHORD
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.systemui.Flags
+import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.screenshot.ImageCapture
 import com.android.systemui.screenshot.ScreenshotData
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.emptyDisplayContent
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.launcherOnly
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
 import com.android.systemui.screenshot.data.repository.DisplayContentRepository
+import com.android.systemui.screenshot.data.repository.profileTypeRepository
+import com.android.systemui.screenshot.policy.CaptureType.FullScreen
+import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
 import com.android.systemui.screenshot.policy.TestUserIds.PERSONAL
 import com.android.systemui.screenshot.policy.TestUserIds.WORK
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
 @RunWith(AndroidJUnit4::class)
 class PolicyRequestProcessorTest {
+    private val kosmos = Kosmos()
+
+    private val screenshotRequest =
+        ScreenshotData(
+            TAKE_SCREENSHOT_FULLSCREEN,
+            SCREENSHOT_KEY_CHORD,
+            UserHandle.CURRENT,
+            topComponent = null,
+            originalScreenBounds = FULL_SCREEN,
+            taskId = -1,
+            originalInsets = Insets.NONE,
+            bitmap = null,
+            displayId = DEFAULT_DISPLAY,
+        )
+
+    val defaultComponent = ComponentName("default", "Component")
+    val defaultOwner = UserHandle.of(PERSONAL)
+
+    @get:Rule val setFlagsRule: SetFlagsRule = SetFlagsRule()
+
+    /** Tests applying CaptureParameters with 'IsolatedTask' CaptureType */
+    @Test
+    @EnableFlags(Flags.FLAG_SCREENSHOT_POLICY_SPLIT_AND_DESKTOP_MODE)
+    fun testProcess_newPolicy_isolatedTask() = runTest {
+        val taskImage = Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(task = taskImage),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
+                policies = emptyList(),
+                defaultOwner = defaultOwner,
+                defaultComponent = defaultComponent,
+                displayTasks = { emptyDisplayContent },
+            )
+
+        val result =
+            requestProcessor.modify(
+                screenshotRequest,
+                CaptureParameters(
+                    IsolatedTask(taskId = TASK_ID, taskBounds = null),
+                    ComponentName.unflattenFromString(FILES),
+                    UserHandle.of(WORK),
+                ),
+            )
+
+        assertWithMessage("The screenshot bitmap").that(result.bitmap).isSameInstanceAs(taskImage)
+
+        assertWithMessage("The assigned owner of the screenshot")
+            .that(result.userHandle)
+            .isEqualTo(UserHandle.of(WORK))
+
+        assertWithMessage("The topComponent of the screenshot")
+            .that(result.topComponent)
+            .isEqualTo(ComponentName.unflattenFromString(FILES))
+
+        assertWithMessage("Task ID").that(result.taskId).isEqualTo(TASK_ID)
+    }
+
+    /** Tests applying CaptureParameters with 'FullScreen' CaptureType */
+    @Test
+    @EnableFlags(Flags.FLAG_SCREENSHOT_POLICY_SPLIT_AND_DESKTOP_MODE)
+    fun testProcess_newPolicy_fullScreen() = runTest {
+        val screenImage = Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(display = screenImage),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
+                policies = emptyList(),
+                defaultOwner = defaultOwner,
+                defaultComponent = defaultComponent,
+                displayTasks = { emptyDisplayContent },
+            )
+
+        val result =
+            requestProcessor.modify(
+                screenshotRequest,
+                CaptureParameters(FullScreen(displayId = 0), defaultComponent, defaultOwner),
+            )
+
+        assertWithMessage("The result bitmap").that(result.bitmap).isSameInstanceAs(screenImage)
+
+        assertWithMessage("The assigned owner of the screenshot")
+            .that(result.userHandle)
+            .isEqualTo(defaultOwner)
+
+        assertWithMessage("The topComponent of the screenshot")
+            .that(result.topComponent)
+            .isEqualTo(defaultComponent)
+
+        assertWithMessage("Task ID").that(result.taskId).isEqualTo(-1)
+    }
+
     /** Tests behavior when no policies are applied */
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_POLICY_SPLIT_AND_DESKTOP_MODE)
     fun testProcess_defaultOwner_whenNoPolicyApplied() {
         val fullScreenWork = DisplayContentRepository {
             singleFullScreen(TaskSpec(taskId = TASK_ID, name = FILES, userId = WORK))
@@ -67,6 +178,7 @@
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = emptyList(),
                 defaultOwner = UserHandle.of(PERSONAL),
                 defaultComponent = ComponentName("default", "Component"),
@@ -95,6 +207,7 @@
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(display = null),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = emptyList(),
                 defaultComponent = ComponentName("default", "Component"),
                 displayTasks = DisplayContentRepository { launcherOnly() },
@@ -118,7 +231,7 @@
                 reason = "",
                 parameters =
                     CaptureParameters(
-                        CaptureType.IsolatedTask(taskId = 0, taskBounds = null),
+                        IsolatedTask(taskId = 0, taskBounds = null),
                         null,
                         UserHandle.CURRENT,
                     ),
@@ -130,6 +243,7 @@
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(task = null),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = listOf(captureTaskPolicy),
                 defaultComponent = ComponentName("default", "Component"),
                 displayTasks = fullScreenWork,