DesktopMode: Make DesktopWallpaperActivity per-display.

This uses a combination of Intent flags - ACTIVITY_NEW_TASK and
FLAG_ACTIVITY_MULTIPLE_TASK and checking TokenProvider's current
existing token to ensure we aren't spawning new task containing
WallpaperActivity if they already exist.

Bug: 381935663
Test: atest and Manual checking on Connected Displays
Flag:com.android.window.flags.enable_per_display_desktop_wallpaper_activity

Change-Id: I95d4106a5e9194ece2931309c9908626c2060b22

4 files changed, 87 insertions, 33 deletions

diff --git a/libs/WindowManager/Shell/AndroidManifest.xml b/libs/WindowManager/Shell/AndroidManifest.xml
index b2ac640a468d..0cfd600a8281 100644
--- a/libs/WindowManager/Shell/AndroidManifest.xml
+++ b/libs/WindowManager/Shell/AndroidManifest.xml
@@ -31,7 +31,7 @@
         <activity
             android:name=".desktopmode.DesktopWallpaperActivity"
             android:excludeFromRecents="true"
-            android:launchMode="singleInstance"
+            android:launchMode="singleInstancePerTask"
             android:showForAllUsers="true"
             android:theme="@style/DesktopWallpaperTheme" />
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 6013648c9806..3c958a2fcbb8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -552,7 +552,11 @@ class DesktopTasksController(
         )
         val wct = WindowContainerTransaction()
         exitSplitIfApplicable(wct, taskInfo)
-        moveHomeTask(wct, toTop = true)
+        if (Flags.enablePerDisplayDesktopWallpaperActivity()) {
+            moveHomeTask(wct, toTop = true, taskInfo.displayId)
+        } else {
+            moveHomeTask(wct, toTop = true)
+        }
         val taskIdToMinimize =
             bringDesktopAppsToFrontBeforeShowingNewTask(taskInfo.displayId, wct, taskInfo.taskId)
         addMoveToDesktopChanges(wct, taskInfo)
@@ -1309,11 +1313,15 @@ class DesktopTasksController(
     ): Int? {
         logV("bringDesktopAppsToFront, newTaskId=%d", newTaskIdInFront)
         // Move home to front, ensures that we go back home when all desktop windows are closed
-        moveHomeTask(wct, toTop = true)
+        if (Flags.enablePerDisplayDesktopWallpaperActivity()) {
+            moveHomeTask(wct, toTop = true, displayId)
+        } else {
+            moveHomeTask(wct, toTop = true)
+        }
 
         // Currently, we only handle the desktop on the default display really.
         if (
-            (displayId == DEFAULT_DISPLAY || Flags.enableBugFixesForSecondaryDisplay()) &&
+            (displayId == DEFAULT_DISPLAY || Flags.enablePerDisplayDesktopWallpaperActivity()) &&
                 ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY.isTrue()
         ) {
             // Add translucent wallpaper activity to show the wallpaper underneath
@@ -1359,9 +1367,13 @@ class DesktopTasksController(
         return taskIdToMinimize
     }
 
-    private fun moveHomeTask(wct: WindowContainerTransaction, toTop: Boolean) {
+    private fun moveHomeTask(
+        wct: WindowContainerTransaction,
+        toTop: Boolean,
+        displayId: Int = DEFAULT_DISPLAY,
+    ) {
         shellTaskOrganizer
-            .getRunningTasks(context.displayId)
+            .getRunningTasks(displayId)
             .firstOrNull { task -> task.activityType == ACTIVITY_TYPE_HOME }
             ?.let { homeTask -> wct.reorder(homeTask.getToken(), /* onTop= */ toTop) }
     }
@@ -1370,12 +1382,19 @@ class DesktopTasksController(
         logV("addWallpaperActivity")
         if (Flags.enableDesktopWallpaperActivityForSystemUser()) {
             val intent = Intent(context, DesktopWallpaperActivity::class.java)
+            if (
+                desktopWallpaperActivityTokenProvider.getToken(displayId) == null &&
+                    Flags.enablePerDisplayDesktopWallpaperActivity()
+            ) {
+                intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
+                intent.addFlags(Intent.FLAG_ACTIVITY_MULTIPLE_TASK)
+            }
             val options =
                 ActivityOptions.makeBasic().apply {
                     launchWindowingMode = WINDOWING_MODE_FULLSCREEN
                     pendingIntentBackgroundActivityStartMode =
                         ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS
-                    if (Flags.enableBugFixesForSecondaryDisplay()) {
+                    if (Flags.enablePerDisplayDesktopWallpaperActivity()) {
                         launchDisplayId = displayId
                     }
                 }
@@ -1391,13 +1410,20 @@ class DesktopTasksController(
             val userHandle = UserHandle.of(userId)
             val userContext = context.createContextAsUser(userHandle, /* flags= */ 0)
             val intent = Intent(userContext, DesktopWallpaperActivity::class.java)
+            if (
+                desktopWallpaperActivityTokenProvider.getToken(displayId) == null &&
+                    Flags.enablePerDisplayDesktopWallpaperActivity()
+            ) {
+                intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
+                intent.addFlags(Intent.FLAG_ACTIVITY_MULTIPLE_TASK)
+            }
             intent.putExtra(Intent.EXTRA_USER_HANDLE, userId)
             val options =
                 ActivityOptions.makeBasic().apply {
                     launchWindowingMode = WINDOWING_MODE_FULLSCREEN
                     pendingIntentBackgroundActivityStartMode =
                         ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS
-                    if (Flags.enableBugFixesForSecondaryDisplay()) {
+                    if (Flags.enablePerDisplayDesktopWallpaperActivity()) {
                         launchDisplayId = displayId
                     }
                 }
@@ -1414,8 +1440,8 @@ class DesktopTasksController(
         }
     }
 
-    private fun removeWallpaperActivity(wct: WindowContainerTransaction) {
-        desktopWallpaperActivityTokenProvider.getToken()?.let { token ->
+    private fun removeWallpaperActivity(wct: WindowContainerTransaction, displayId: Int) {
+        desktopWallpaperActivityTokenProvider.getToken(displayId)?.let { token ->
             logV("removeWallpaperActivity")
             if (Flags.enableDesktopWallpaperActivityForSystemUser()) {
                 wct.reorder(token, /* onTop= */ false)
@@ -1440,9 +1466,6 @@ class DesktopTasksController(
             if (!taskRepository.isOnlyVisibleNonClosingTask(taskId, displayId)) {
                 return
             }
-            if (displayId != DEFAULT_DISPLAY) {
-                return
-            }
         } else if (
             Flags.enableDesktopWindowingPip() &&
                 taskRepository.isMinimizedPipPresentInDisplay(displayId) &&
@@ -1457,7 +1480,7 @@ class DesktopTasksController(
         desktopModeEnterExitTransitionListener?.onExitDesktopModeTransitionStarted(
             FULLSCREEN_ANIMATION_DURATION
         )
-        removeWallpaperActivity(wct)
+        removeWallpaperActivity(wct, displayId)
     }
 
     fun releaseVisualIndicator() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
index d61ffdaf5cf8..77af809a9679 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
@@ -60,7 +60,9 @@ class DesktopTasksTransitionObserver(
     shellInit: ShellInit,
 ) : Transitions.TransitionObserver {
 
-    private var transitionToCloseWallpaper: IBinder? = null
+    data class CloseWallpaperTransition(val transition: IBinder, val displayId: Int)
+
+    private var transitionToCloseWallpaper: CloseWallpaperTransition? = null
     /* Pending PiP transition and its associated display id and task id. */
     private var pendingPipTransitionAndPipTask: Triple<IBinder, Int, Int>? = null
     private var currentProfileId: Int
@@ -248,9 +250,10 @@ class DesktopTasksTransitionObserver(
                 desktopRepository.getVisibleTaskCount(taskInfo.displayId) == 0 &&
                     change.mode == TRANSIT_CLOSE &&
                     taskInfo.windowingMode == WINDOWING_MODE_FREEFORM &&
-                    desktopWallpaperActivityTokenProvider.getToken() != null
+                    desktopWallpaperActivityTokenProvider.getToken(taskInfo.displayId) != null
             ) {
-                transitionToCloseWallpaper = transition
+                transitionToCloseWallpaper =
+                    CloseWallpaperTransition(transition, taskInfo.displayId)
                 currentProfileId = taskInfo.userId
             }
         }
@@ -265,25 +268,28 @@ class DesktopTasksTransitionObserver(
     }
 
     override fun onTransitionFinished(transition: IBinder, aborted: Boolean) {
+        val lastSeenTransitionToCloseWallpaper = transitionToCloseWallpaper
         // TODO: b/332682201 Update repository state
-        if (transitionToCloseWallpaper == transition) {
+        if (lastSeenTransitionToCloseWallpaper?.transition == transition) {
             // TODO: b/362469671 - Handle merging the animation when desktop is also closing.
-            desktopWallpaperActivityTokenProvider.getToken()?.let { wallpaperActivityToken ->
-                if (Flags.enableDesktopWallpaperActivityForSystemUser()) {
-                    transitions.startTransition(
-                        TRANSIT_TO_BACK,
-                        WindowContainerTransaction()
-                            .reorder(wallpaperActivityToken, /* onTop= */ false),
-                        null,
-                    )
-                } else {
-                    transitions.startTransition(
-                        TRANSIT_CLOSE,
-                        WindowContainerTransaction().removeTask(wallpaperActivityToken),
-                        null,
-                    )
+            desktopWallpaperActivityTokenProvider
+                .getToken(lastSeenTransitionToCloseWallpaper.displayId)
+                ?.let { wallpaperActivityToken ->
+                    if (Flags.enableDesktopWallpaperActivityForSystemUser()) {
+                        transitions.startTransition(
+                            TRANSIT_TO_BACK,
+                            WindowContainerTransaction()
+                                .reorder(wallpaperActivityToken, /* onTop= */ false),
+                            null,
+                        )
+                    } else {
+                        transitions.startTransition(
+                            TRANSIT_CLOSE,
+                            WindowContainerTransaction().removeTask(wallpaperActivityToken),
+                            null,
+                        )
+                    }
                 }
-            }
             transitionToCloseWallpaper = null
         } else if (pendingPipTransitionAndPipTask?.first == transition) {
             val desktopRepository = desktopUserRepositories.getProfile(currentProfileId)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index 4bb743079861..7ea7203e90b0 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -602,7 +602,32 @@ class DesktopTasksControllerTest : ShellTestCase() {
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_PER_DISPLAY_DESKTOP_WALLPAPER_ACTIVITY,
+    )
+    fun showDesktopApps_onSecondaryDisplay_desktopWallpaperEnabled_perDisplayWallpaperEnabled_shouldShowWallpaper() {
+        val homeTask = setUpHomeTask(SECOND_DISPLAY)
+        val task1 = setUpFreeformTask(SECOND_DISPLAY)
+        val task2 = setUpFreeformTask(SECOND_DISPLAY)
+        markTaskHidden(task1)
+        markTaskHidden(task2)
+
+        controller.showDesktopApps(SECOND_DISPLAY, RemoteTransition(TestRemoteTransition()))
+
+        val wct =
+            getLatestWct(type = TRANSIT_TO_FRONT, handlerClass = OneShotRemoteHandler::class.java)
+        assertThat(wct.hierarchyOps).hasSize(4)
+        // Expect order to be from bottom: home, wallpaperIntent, task1, task2
+        wct.assertReorderAt(index = 0, homeTask)
+        wct.assertPendingIntentAt(index = 1, desktopWallpaperIntent)
+        wct.assertReorderAt(index = 2, task1)
+        wct.assertReorderAt(index = 3, task2)
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    @DisableFlags(Flags.FLAG_ENABLE_PER_DISPLAY_DESKTOP_WALLPAPER_ACTIVITY)
     fun showDesktopApps_onSecondaryDisplay_desktopWallpaperEnabled_shouldNotShowWallpaper() {
         val homeTask = setUpHomeTask(SECOND_DISPLAY)
         val task1 = setUpFreeformTask(SECOND_DISPLAY)