Improved screenshot policy for desktop and split mode

The new unified policy class simplifies the rules and behavior
of screenshots, handling situations with mixed profile content
on screen.

Improvements:

When two work profile tasks apear as split-screen mode, the
screenshot will now capture both.

When desktop mode is active with freeform windows, content
will be saved to the private profile if any is private profile,
otherwise to the personal profile.

When in desktop mode with a single maximized window, it is
treated as a full screen task and handled accordingly.

The focused task is no longer used to determine the owner of the
screenshot making the result less surprising to end users.

If work content appears beside personal content in split mode,
the screenshot will remain full-screen and be saved to the personal
profile.

Bug: 365597999
Flag: com.android.systemui.screenshot_policy_split_and_desktop_mode
Test: atest ScreenshotPolicyTest
Change-Id: I5353b477d527b79f8b95803d7ba6136fdba441be

12 files changed, 798 insertions, 147 deletions

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 254f1e1efe13..4d71dc45001d 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 android.graphics.Rect
 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 @@ object DisplayContentScenarios {
 
     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 @@ object DisplayContentScenarios {
             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 @@ object DisplayContentScenarios {
                             .sortedBy { it.id == focusedTaskId }
                     },
                     launcher(visible = false),
-                )
+                ),
         )
     }
 
@@ -124,7 +146,7 @@ object DisplayContentScenarios {
                     fullScreen?.also { add(fullScreen(it, visible = true)) }
                     add(launcher(visible = (fullScreen == null)))
                     add(emptyRootSplit)
-                }
+                },
         )
     }
 
@@ -142,7 +164,7 @@ object DisplayContentScenarios {
         return DisplayContentModel(
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
-            rootTasks = freeFormTasks + launcher(visible = true) + emptyRootSplit
+            rootTasks = freeFormTasks + launcher(visible = true) + emptyRootSplit,
         )
     }
 
@@ -153,11 +175,18 @@ object DisplayContentScenarios {
      * 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 @@ object DisplayContentScenarios {
             "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 @@ object DisplayContentScenarios {
             }
 
         /** 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 @@ object DisplayContentScenarios {
                 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 6c35b233ffec..cedf0c8a2c06 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 @@ fun RootTaskInfo.newChildTask(
     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 @@ fun newRootTaskInfo(
     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 6e57761aafa6..b7f565df4a3c 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 @@ class PrivateProfilePolicyTest {
             policy.check(
                 singleFullScreen(
                     spec = TaskSpec(taskId = 1002, name = YOUTUBE, userId = PRIVATE),
-                    shadeExpanded = true
+                    shadeExpanded = true,
                 )
             )
 
@@ -93,8 +93,8 @@ class PrivateProfilePolicyTest {
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -110,25 +110,20 @@ class PrivateProfilePolicyTest {
                         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 @@ class PrivateProfilePolicyTest {
         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 @@ class PrivateProfilePolicyTest {
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -161,9 +156,9 @@ class PrivateProfilePolicyTest {
         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 @@ class PrivateProfilePolicyTest {
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(FILES),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -196,8 +191,8 @@ class PrivateProfilePolicyTest {
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE_PIP),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -220,8 +215,8 @@ class PrivateProfilePolicyTest {
                     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 000000000000..28eb9fc1364b
--- /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 be9fcc2be3a3..30a786c291cb 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.Activi
 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 @@ class WorkProfilePolicyTest {
     @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 @@ class WorkProfilePolicyTest {
                 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 @@ class WorkProfilePolicyTest {
                 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 @@ class WorkProfilePolicyTest {
             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 @@ class WorkProfilePolicyTest {
                         type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -166,9 +149,11 @@ class WorkProfilePolicyTest {
         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 @@ class WorkProfilePolicyTest {
                     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 @@ class WorkProfilePolicyTest {
         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 @@ class WorkProfilePolicyTest {
                         type = IsolatedTask(taskId = 1003, taskBounds = FULL_SCREEN),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -238,7 +217,7 @@ class WorkProfilePolicyTest {
                 freeFormApps(
                     TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
                     TaskSpec(taskId = 1003, name = FILES, userId = WORK),
-                    focusedTaskId = 1003
+                    focusedTaskId = 1003,
                 )
             )
 
@@ -251,7 +230,7 @@ class WorkProfilePolicyTest {
                         type = IsolatedTask(taskId = 1003, taskBounds = FREE_FORM),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -264,16 +243,10 @@ class WorkProfilePolicyTest {
                 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 0ef5207cad37..945520126474 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 @@ sealed interface CaptureType {
     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 039143a1907d..e840668688a0 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.os.UserHandle
 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 @@ class PolicyRequestProcessor(
     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 @@ class PolicyRequestProcessor(
         }
         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 @@ class PolicyRequestProcessor(
     }
 
     /** 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 @@ class PolicyRequestProcessor(
                         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 f768cfb2ceb5..dd39f92643ce 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 @@ internal fun RootTaskInfo.childTasksTopDown(): Sequence<ChildTaskModel> {
             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 000000000000..9967afffb6a0
--- /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 2cb9fe7f1a9d..a9c6370bb776 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 @@ import kotlinx.coroutines.CoroutineDispatcher
 
 @Module
 interface ScreenshotPolicyModule {
-
     @Binds
     @SysUISingleton
     fun bindProfileTypeRepository(impl: ProfileTypeRepositoryImpl): ProfileTypeRepository
@@ -67,6 +66,7 @@ interface ScreenshotPolicyModule {
             imageCapture: ImageCapture,
             displayContentRepo: DisplayContentRepository,
             policyListProvider: Provider<List<CapturePolicy>>,
+            standardPolicy: ScreenshotPolicy,
         ): ScreenshotRequestProcessor {
             return PolicyRequestProcessor(
                 background = background,
@@ -75,7 +75,8 @@ interface ScreenshotPolicyModule {
                 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 29450a20b1d3..cf90c0a58e94 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.CapturePolicy.PolicyResult.NotMatc
 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 @@ import kotlinx.coroutines.flow.first
  */
 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 @@ constructor(
                 .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 @@ constructor(
                 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 0d4cb4c6751c..7709a65712a1 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.Bitmap
 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 @@ class PolicyRequestProcessorTest {
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = emptyList(),
                 defaultOwner = UserHandle.of(PERSONAL),
                 defaultComponent = ComponentName("default", "Component"),
@@ -95,6 +207,7 @@ class PolicyRequestProcessorTest {
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(display = null),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = emptyList(),
                 defaultComponent = ComponentName("default", "Component"),
                 displayTasks = DisplayContentRepository { launcherOnly() },
@@ -118,7 +231,7 @@ class PolicyRequestProcessorTest {
                 reason = "",
                 parameters =
                     CaptureParameters(
-                        CaptureType.IsolatedTask(taskId = 0, taskBounds = null),
+                        IsolatedTask(taskId = 0, taskBounds = null),
                         null,
                         UserHandle.CURRENT,
                     ),
@@ -130,6 +243,7 @@ class PolicyRequestProcessorTest {
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(task = null),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = listOf(captureTaskPolicy),
                 defaultComponent = ComponentName("default", "Component"),
                 displayTasks = fullScreenWork,