Merge "Refactor wallpaper activity token outside of repo." into main

7 files changed, 90 insertions, 80 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index ace7f078bb10..d02c6b05e5b6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -100,6 +100,7 @@ import com.android.wm.shell.desktopmode.SpringDragToDesktopTransitionHandler;
 import com.android.wm.shell.desktopmode.ToggleResizeDesktopTaskTransitionHandler;
 import com.android.wm.shell.desktopmode.WindowDecorCaptionHandleRepository;
 import com.android.wm.shell.desktopmode.compatui.SystemModalsTransitionHandler;
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider;
 import com.android.wm.shell.desktopmode.education.AppHandleEducationController;
 import com.android.wm.shell.desktopmode.education.AppHandleEducationFilter;
 import com.android.wm.shell.desktopmode.education.AppToWebEducationController;
@@ -733,7 +734,8 @@ public abstract class WMShellModule {
             FocusTransitionObserver focusTransitionObserver,
             DesktopModeEventLogger desktopModeEventLogger,
             DesktopModeUiEventLogger desktopModeUiEventLogger,
-            DesktopTilingDecorViewModel desktopTilingDecorViewModel) {
+            DesktopTilingDecorViewModel desktopTilingDecorViewModel,
+            DesktopWallpaperActivityTokenProvider desktopWallpaperActivityTokenProvider) {
         return new DesktopTasksController(
                 context,
                 shellInit,
@@ -764,7 +766,8 @@ public abstract class WMShellModule {
                 mainHandler,
                 desktopModeEventLogger,
                 desktopModeUiEventLogger,
-                desktopTilingDecorViewModel);
+                desktopTilingDecorViewModel,
+                desktopWallpaperActivityTokenProvider);
     }
 
     @WMSingleton
@@ -1092,6 +1095,7 @@ public abstract class WMShellModule {
             ShellTaskOrganizer shellTaskOrganizer,
             Optional<DesktopMixedTransitionHandler> desktopMixedTransitionHandler,
             Optional<BackAnimationController> backAnimationController,
+            DesktopWallpaperActivityTokenProvider desktopWallpaperActivityTokenProvider,
             ShellInit shellInit) {
         return desktopUserRepositories.flatMap(
                 repository ->
@@ -1103,6 +1107,7 @@ public abstract class WMShellModule {
                                         shellTaskOrganizer,
                                         desktopMixedTransitionHandler.get(),
                                         backAnimationController.get(),
+                                        desktopWallpaperActivityTokenProvider,
                                         shellInit)));
     }
 
@@ -1306,6 +1311,12 @@ public abstract class WMShellModule {
         return new DesktopModeUiEventLogger(uiEventLogger, packageManager);
     }
 
+    @WMSingleton
+    @Provides
+    static DesktopWallpaperActivityTokenProvider provideDesktopWallpaperActivityTokenProvider() {
+        return new DesktopWallpaperActivityTokenProvider();
+    }
+
     //
     // Drag and drop
     //
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
index c5b570dd3d57..d3066645f32e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
@@ -23,7 +23,6 @@ import android.util.ArraySet
 import android.util.SparseArray
 import android.view.Display.INVALID_DISPLAY
 import android.window.DesktopModeFlags
-import android.window.WindowContainerToken
 import androidx.core.util.forEach
 import androidx.core.util.keyIterator
 import androidx.core.util.valueIterator
@@ -90,9 +89,6 @@ class DesktopRepository(
         }
     }
 
-    /* Current wallpaper activity token to remove wallpaper activity when last task is removed. */
-    var wallpaperActivityToken: WindowContainerToken? = null
-
     private val activeTasksListeners = ArraySet<ActiveTasksListener>()
     private val visibleTasksListeners = ArrayMap<VisibleTasksListener, Executor>()
 
@@ -549,7 +545,6 @@ class DesktopRepository(
                 "${innerPrefix}topTransparentFullscreenTaskId=" +
                     "${data.topTransparentFullscreenTaskId}"
             )
-            pw.println("${innerPrefix}wallpaperActivityToken=$wallpaperActivityToken")
         }
     }
 
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 e43ad52ef105..1af127470327 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
@@ -95,6 +95,7 @@ import com.android.wm.shell.desktopmode.DragToDesktopTransitionHandler.DragToDes
 import com.android.wm.shell.desktopmode.EnterDesktopTaskTransitionHandler.FREEFORM_ANIMATION_DURATION
 import com.android.wm.shell.desktopmode.ExitDesktopTaskTransitionHandler.FULLSCREEN_ANIMATION_DURATION
 import com.android.wm.shell.desktopmode.common.ToggleTaskSizeInteraction
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider
 import com.android.wm.shell.desktopmode.minimize.DesktopWindowLimitRemoteHandler
 import com.android.wm.shell.draganddrop.DragAndDropController
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter
@@ -170,6 +171,7 @@ class DesktopTasksController(
     private val desktopModeEventLogger: DesktopModeEventLogger,
     private val desktopModeUiEventLogger: DesktopModeUiEventLogger,
     private val desktopTilingDecorViewModel: DesktopTilingDecorViewModel,
+    private val desktopWallpaperActivityTokenProvider: DesktopWallpaperActivityTokenProvider,
 ) :
     RemoteCallable<DesktopTasksController>,
     Transitions.TransitionHandler,
@@ -1366,7 +1368,7 @@ class DesktopTasksController(
     }
 
     private fun removeWallpaperActivity(wct: WindowContainerTransaction) {
-        taskRepository.wallpaperActivityToken?.let { token ->
+        desktopWallpaperActivityTokenProvider.getToken()?.let { token ->
             logV("removeWallpaperActivity")
             wct.removeTask(token)
         }
@@ -1396,9 +1398,7 @@ class DesktopTasksController(
         desktopModeEnterExitTransitionListener?.onExitDesktopModeTransitionStarted(
             FULLSCREEN_ANIMATION_DURATION
         )
-        if (taskRepository.wallpaperActivityToken != null) {
-            removeWallpaperActivity(wct)
-        }
+        removeWallpaperActivity(wct)
     }
 
     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 5c79658b6809..2bb299cc1fb9 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
@@ -32,6 +32,7 @@ import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.back.BackAnimationController
 import com.android.wm.shell.desktopmode.DesktopModeTransitionTypes.isExitDesktopModeTransition
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
@@ -50,6 +51,7 @@ class DesktopTasksTransitionObserver(
     private val shellTaskOrganizer: ShellTaskOrganizer,
     private val desktopMixedTransitionHandler: DesktopMixedTransitionHandler,
     private val backAnimationController: BackAnimationController,
+    private val desktopWallpaperActivityTokenProvider: DesktopWallpaperActivityTokenProvider,
     shellInit: ShellInit,
 ) : Transitions.TransitionObserver {
 
@@ -212,7 +214,7 @@ class DesktopTasksTransitionObserver(
                 desktopRepository.getVisibleTaskCount(taskInfo.displayId) == 0 &&
                     change.mode == TRANSIT_CLOSE &&
                     taskInfo.windowingMode == WINDOWING_MODE_FREEFORM &&
-                    desktopRepository.wallpaperActivityToken != null
+                    desktopWallpaperActivityTokenProvider.getToken() != null
             ) {
                 transitionToCloseWallpaper = transition
                 currentProfileId = taskInfo.userId
@@ -232,8 +234,7 @@ class DesktopTasksTransitionObserver(
         // TODO: b/332682201 Update repository state
         if (transitionToCloseWallpaper == transition) {
             // TODO: b/362469671 - Handle merging the animation when desktop is also closing.
-            val desktopRepository = desktopUserRepositories.getProfile(currentProfileId)
-            desktopRepository.wallpaperActivityToken?.let { wallpaperActivityToken ->
+            desktopWallpaperActivityTokenProvider.getToken()?.let { wallpaperActivityToken ->
                 transitions.startTransition(
                     TRANSIT_CLOSE,
                     WindowContainerTransaction().removeTask(wallpaperActivityToken),
@@ -251,10 +252,12 @@ class DesktopTasksTransitionObserver(
         info.changes.forEach { change ->
             change.taskInfo?.let { taskInfo ->
                 if (DesktopWallpaperActivity.isWallpaperTask(taskInfo)) {
-                    val desktopRepository = desktopUserRepositories.getProfile(taskInfo.userId)
                     when (change.mode) {
                         WindowManager.TRANSIT_OPEN -> {
-                            desktopRepository.wallpaperActivityToken = taskInfo.token
+                            desktopWallpaperActivityTokenProvider.setToken(
+                                taskInfo.token,
+                                taskInfo.displayId,
+                            )
                             // After the task for the wallpaper is created, set it non-trimmable.
                             // This is important to prevent recents from trimming and removing the
                             // task.
@@ -263,7 +266,8 @@ class DesktopTasksTransitionObserver(
                                     .setTaskTrimmableFromRecents(taskInfo.token, false)
                             )
                         }
-                        TRANSIT_CLOSE -> desktopRepository.wallpaperActivityToken = null
+                        TRANSIT_CLOSE ->
+                            desktopWallpaperActivityTokenProvider.removeToken(taskInfo.displayId)
                         else -> {}
                     }
                 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/desktopwallpaperactivity/DesktopWallpaperActivityTokenProvider.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/desktopwallpaperactivity/DesktopWallpaperActivityTokenProvider.kt
new file mode 100644
index 000000000000..a87004c07d43
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/desktopwallpaperactivity/DesktopWallpaperActivityTokenProvider.kt
@@ -0,0 +1,39 @@
+/*
+ * 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.wm.shell.desktopmode.desktopwallpaperactivity
+
+import android.util.SparseArray
+import android.view.Display.DEFAULT_DISPLAY
+import android.window.WindowContainerToken
+
+/** Provides per display window container tokens for [DesktopWallpaperActivity]. */
+class DesktopWallpaperActivityTokenProvider {
+
+    private val wallpaperActivityTokenByDisplayId = SparseArray<WindowContainerToken>()
+
+    fun setToken(token: WindowContainerToken, displayId: Int = DEFAULT_DISPLAY) {
+        wallpaperActivityTokenByDisplayId[displayId] = token
+    }
+
+    fun getToken(displayId: Int = DEFAULT_DISPLAY): WindowContainerToken? {
+        return wallpaperActivityTokenByDisplayId[displayId]
+    }
+
+    fun removeToken(displayId: Int = DEFAULT_DISPLAY) {
+        wallpaperActivityTokenByDisplayId.delete(displayId)
+    }
+}
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 61aec739724a..45eef2a85f55 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
@@ -111,6 +111,7 @@ import com.android.wm.shell.desktopmode.DesktopTestHelpers.createSplitScreenTask
 import com.android.wm.shell.desktopmode.EnterDesktopTaskTransitionHandler.FREEFORM_ANIMATION_DURATION
 import com.android.wm.shell.desktopmode.ExitDesktopTaskTransitionHandler.FULLSCREEN_ANIMATION_DURATION
 import com.android.wm.shell.desktopmode.common.ToggleTaskSizeInteraction
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider
 import com.android.wm.shell.desktopmode.minimize.DesktopWindowLimitRemoteHandler
 import com.android.wm.shell.desktopmode.persistence.Desktop
 import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
@@ -236,6 +237,10 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @Mock
     lateinit var desktopModeEnterExitTransitionListener: DesktopModeEntryExitTransitionListener
     @Mock private lateinit var userManager: UserManager
+    @Mock
+    private lateinit var desktopWallpaperActivityTokenProvider:
+        DesktopWallpaperActivityTokenProvider
+
     private lateinit var controller: DesktopTasksController
     private lateinit var shellInit: ShellInit
     private lateinit var taskRepository: DesktopRepository
@@ -257,6 +262,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     private val RESIZABLE_PORTRAIT_BOUNDS = Rect(680, 75, 1880, 1275)
     private val UNRESIZABLE_LANDSCAPE_BOUNDS = Rect(25, 449, 1575, 1611)
     private val UNRESIZABLE_PORTRAIT_BOUNDS = Rect(830, 75, 1730, 1275)
+    private val wallpaperToken = MockToken().token()
 
     @Before
     fun setUp() {
@@ -324,6 +330,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
                 )
             )
             .thenReturn(ExitResult.NoExit)
+        whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(wallpaperToken)
 
         controller = createController()
         controller.setSplitScreenController(splitScreenController)
@@ -375,6 +382,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
             desktopModeEventLogger,
             desktopModeUiEventLogger,
             desktopTilingDecorViewModel,
+            desktopWallpaperActivityTokenProvider,
         )
     }
 
@@ -1488,9 +1496,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @Test
     fun moveToFullscreen_tdaFullscreen_windowingModeUndefined_removesWallpaperActivity() {
         val task = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
-
-        taskRepository.wallpaperActivityToken = wallpaperToken
         assertNotNull(rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY))
             .configuration
             .windowConfiguration
@@ -1523,9 +1528,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @Test
     fun moveToFullscreen_tdaFreeform_windowingModeFullscreen_removesWallpaperActivity() {
         val task = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         assertNotNull(rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY))
             .configuration
             .windowConfiguration
@@ -1547,9 +1550,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task1 = setUpFreeformTask()
         // Setup task2
         setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         assertNotNull(rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY))
             .configuration
             .windowConfiguration
@@ -1748,8 +1749,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
             .thenReturn(secondDisplayArea)
         // Add a task and a wallpaper
         val task = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
 
         controller.moveToNextDisplay(task.taskId)
 
@@ -1960,7 +1959,10 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun onDesktopWindowClose_singleActiveTask_noWallpaperActivityToken() {
         val task = setUpFreeformTask()
         val wct = WindowContainerTransaction()
+        whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
+
         controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task)
+
         // Doesn't modify transaction
         assertThat(wct.hierarchyOps).isEmpty()
     }
@@ -1968,8 +1970,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @Test
     fun onDesktopWindowClose_singleActiveTask_hasWallpaperActivityToken() {
         val task = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
 
         val wct = WindowContainerTransaction()
         controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task)
@@ -1980,8 +1980,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @Test
     fun onDesktopWindowClose_singleActiveTask_isClosing() {
         val task = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
+
         taskRepository.addClosingTask(DEFAULT_DISPLAY, task.taskId)
 
         val wct = WindowContainerTransaction()
@@ -1993,8 +1992,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @Test
     fun onDesktopWindowClose_singleActiveTask_isMinimized() {
         val task = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
+
         taskRepository.minimizeTask(DEFAULT_DISPLAY, task.taskId)
 
         val wct = WindowContainerTransaction()
@@ -2007,8 +2005,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun onDesktopWindowClose_multipleActiveTasks() {
         val task1 = setUpFreeformTask()
         setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
 
         val wct = WindowContainerTransaction()
         controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task1)
@@ -2020,8 +2016,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun onDesktopWindowClose_multipleActiveTasks_isOnlyNonClosingTask() {
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
+
         taskRepository.addClosingTask(DEFAULT_DISPLAY, task2.taskId)
 
         val wct = WindowContainerTransaction()
@@ -2034,8 +2029,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun onDesktopWindowClose_multipleActiveTasks_hasMinimized() {
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
+
         taskRepository.minimizeTask(DEFAULT_DISPLAY, task2.taskId)
 
         val wct = WindowContainerTransaction()
@@ -2050,8 +2044,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val transition = Binder()
         whenever(freeformTaskTransitionStarter.startMinimizedModeTransition(any()))
             .thenReturn(transition)
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
 
         controller.minimizeTask(task)
 
@@ -2108,8 +2100,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val transition = Binder()
         whenever(freeformTaskTransitionStarter.startMinimizedModeTransition(any()))
             .thenReturn(transition)
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
 
         // The only active task is being minimized.
         controller.minimizeTask(task)
@@ -2126,8 +2116,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val transition = Binder()
         whenever(freeformTaskTransitionStarter.startMinimizedModeTransition(any()))
             .thenReturn(transition)
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.minimizeTask(DEFAULT_DISPLAY, task.taskId)
 
         // The only active task is already minimized.
@@ -2147,8 +2135,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val transition = Binder()
         whenever(freeformTaskTransitionStarter.startMinimizedModeTransition(any()))
             .thenReturn(transition)
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
 
         controller.minimizeTask(task1)
 
@@ -2166,8 +2152,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val transition = Binder()
         whenever(freeformTaskTransitionStarter.startMinimizedModeTransition(any()))
             .thenReturn(transition)
-        val wallpaperToken = MockToken().token()
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.minimizeTask(DEFAULT_DISPLAY, task2.taskId)
 
         // task1 is the only visible task as task2 is minimized.
@@ -2776,6 +2760,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun handleRequest_backTransition_singleTaskNoToken_withWallpaper_removesTask() {
         val task = setUpFreeformTask()
+        whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
 
         val result =
             controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
@@ -2799,6 +2784,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun handleRequest_backTransition_singleTaskNoToken_doesNotHandle() {
         val task = setUpFreeformTask()
+        whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
 
         val result =
             controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
@@ -2811,7 +2797,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun handleRequest_backTransition_singleTaskWithToken_noWallpaper_doesNotHandle() {
         val task = setUpFreeformTask()
 
-        taskRepository.wallpaperActivityToken = MockToken().token()
         val result =
             controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
 
@@ -2822,9 +2807,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun handleRequest_backTransition_singleTaskWithToken_removesWallpaper() {
         val task = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         val result =
             controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
 
@@ -2838,7 +2821,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task1 = setUpFreeformTask()
         setUpFreeformTask()
 
-        taskRepository.wallpaperActivityToken = MockToken().token()
         val result =
             controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
 
@@ -2851,7 +2833,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task1 = setUpFreeformTask()
         setUpFreeformTask()
 
-        taskRepository.wallpaperActivityToken = MockToken().token()
         val result =
             controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
 
@@ -2866,9 +2847,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun handleRequest_backTransition_multipleTasksSingleNonClosing_removesWallpaperAndTask() {
         val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
         val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.addClosingTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
         val result =
             controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
@@ -2882,9 +2861,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun handleRequest_backTransition_multipleTasksSingleNonMinimized_removesWallpaperAndTask() {
         val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
         val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.minimizeTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
         val result =
             controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
@@ -2898,9 +2875,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun handleRequest_backTransition_nonMinimizadTask_withWallpaper_removesWallpaper() {
         val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
         val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.minimizeTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
         // Task is being minimized so mark it as not visible.
         taskRepository.updateTask(displayId = DEFAULT_DISPLAY, task2.taskId, isVisible = false)
@@ -2925,6 +2900,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun handleRequest_closeTransition_singleTaskNoToken_doesNotHandle() {
         val task = setUpFreeformTask()
+        whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
 
         val result =
             controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_CLOSE))
@@ -2937,7 +2913,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun handleRequest_closeTransition_singleTaskWithToken_noWallpaper_doesNotHandle() {
         val task = setUpFreeformTask()
 
-        taskRepository.wallpaperActivityToken = MockToken().token()
         val result =
             controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_CLOSE))
 
@@ -2948,9 +2923,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun handleRequest_closeTransition_singleTaskWithToken_withWallpaper_removesWallpaper() {
         val task = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         val result =
             controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_CLOSE))
 
@@ -2964,7 +2937,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task1 = setUpFreeformTask()
         setUpFreeformTask()
 
-        taskRepository.wallpaperActivityToken = MockToken().token()
         val result =
             controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_CLOSE))
 
@@ -2977,7 +2949,6 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task1 = setUpFreeformTask()
         setUpFreeformTask()
 
-        taskRepository.wallpaperActivityToken = MockToken().token()
         val result =
             controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_CLOSE))
 
@@ -2989,9 +2960,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun handleRequest_closeTransition_multipleTasksSingleNonClosing_removesWallpaper() {
         val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
         val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.addClosingTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
         val result =
             controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_CLOSE))
@@ -3005,9 +2974,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun handleRequest_closeTransition_multipleTasksSingleNonMinimized_removesWallpaper() {
         val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
         val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.minimizeTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
         val result =
             controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_CLOSE))
@@ -3021,9 +2988,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     fun handleRequest_closeTransition_minimizadTask_withWallpaper_removesWallpaper() {
         val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
         val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
-        val wallpaperToken = MockToken().token()
 
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.minimizeTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
         // Task is being minimized so mark it as not visible.
         taskRepository.updateTask(displayId = DEFAULT_DISPLAY, task2.taskId, isVisible = false)
@@ -3099,12 +3064,10 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
         val task3 = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
 
         task1.isFocused = false
         task2.isFocused = true
         task3.isFocused = false
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.minimizeTask(DEFAULT_DISPLAY, task1.taskId)
         taskRepository.updateTask(DEFAULT_DISPLAY, task3.taskId, isVisible = false)
 
@@ -3122,12 +3085,10 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
         val task3 = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
 
         task1.isFocused = false
         task2.isFocused = true
         task3.isFocused = false
-        taskRepository.wallpaperActivityToken = wallpaperToken
         controller.enterFullscreen(DEFAULT_DISPLAY, transitionSource = UNKNOWN)
 
         val wct = getLatestExitDesktopWct()
@@ -3614,12 +3575,10 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
         val task3 = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
 
         task1.isFocused = false
         task2.isFocused = true
         task3.isFocused = false
-        taskRepository.wallpaperActivityToken = wallpaperToken
         taskRepository.minimizeTask(DEFAULT_DISPLAY, task1.taskId)
         taskRepository.updateTask(DEFAULT_DISPLAY, task3.taskId, isVisible = false)
 
@@ -3642,12 +3601,10 @@ class DesktopTasksControllerTest : ShellTestCase() {
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
         val task3 = setUpFreeformTask()
-        val wallpaperToken = MockToken().token()
 
         task1.isFocused = false
         task2.isFocused = true
         task3.isFocused = false
-        taskRepository.wallpaperActivityToken = wallpaperToken
 
         controller.enterSplit(DEFAULT_DISPLAY, leftOrTop = false)
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt
index c66d203fd89a..176102a4deff 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt
@@ -43,6 +43,7 @@ import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.back.BackAnimationController
 import com.android.wm.shell.common.ShellExecutor
 import com.android.wm.shell.desktopmode.DesktopModeTransitionTypes.TRANSIT_EXIT_DESKTOP_MODE_TASK_DRAG
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
@@ -87,6 +88,9 @@ class DesktopTasksTransitionObserverTest {
     private val taskRepository = mock<DesktopRepository>()
     private val mixedHandler = mock<DesktopMixedTransitionHandler>()
     private val backAnimationController = mock<BackAnimationController>()
+    private val desktopWallpaperActivityTokenProvider =
+        mock<DesktopWallpaperActivityTokenProvider>()
+    private val wallpaperToken = MockToken().token()
 
     private lateinit var transitionObserver: DesktopTasksTransitionObserver
     private lateinit var shellInit: ShellInit
@@ -98,6 +102,7 @@ class DesktopTasksTransitionObserverTest {
 
         whenever(userRepositories.current).thenReturn(taskRepository)
         whenever(userRepositories.getProfile(anyInt())).thenReturn(taskRepository)
+        whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(wallpaperToken)
 
         transitionObserver =
             DesktopTasksTransitionObserver(
@@ -107,6 +112,7 @@ class DesktopTasksTransitionObserverTest {
                 shellTaskOrganizer,
                 mixedHandler,
                 backAnimationController,
+                desktopWallpaperActivityTokenProvider,
                 shellInit,
             )
     }
@@ -236,9 +242,7 @@ class DesktopTasksTransitionObserverTest {
     fun closeLastTask_wallpaperTokenExists_wallpaperIsRemoved() {
         val mockTransition = Mockito.mock(IBinder::class.java)
         val task = createTaskInfo(1, WINDOWING_MODE_FREEFORM)
-        val wallpaperToken = MockToken().token()
         whenever(taskRepository.getVisibleTaskCount(task.displayId)).thenReturn(0)
-        whenever(taskRepository.wallpaperActivityToken).thenReturn(wallpaperToken)
 
         transitionObserver.onTransitionReady(
             transition = mockTransition,