Merge "Remove DesktopWallpaperActivity token onVanished." into main

3 files changed, 44 insertions, 4 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index 43f1a1037cab..0e6481b1c0ac 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -581,12 +581,15 @@ public abstract class WMShellBaseModule {
             ShellTaskOrganizer shellTaskOrganizer,
             SyncTransactionQueue syncQueue,
             Optional<RecentTasksController> recentTasksOptional,
-            Optional<WindowDecorViewModel> windowDecorViewModelOptional) {
+            Optional<WindowDecorViewModel> windowDecorViewModelOptional,
+            Optional<DesktopWallpaperActivityTokenProvider>
+                    desktopWallpaperActivityTokenProviderOptional) {
         if (fullscreenTaskListener.isPresent()) {
             return fullscreenTaskListener.get();
         } else {
             return new FullscreenTaskListener(shellInit, shellTaskOrganizer, syncQueue,
-                    recentTasksOptional, windowDecorViewModelOptional);
+                    recentTasksOptional, windowDecorViewModelOptional,
+                    desktopWallpaperActivityTokenProviderOptional);
         }
     }
 
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
index 2bd7a9873a5e..2a8a3475c2a5 100644
--- 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
@@ -20,6 +20,9 @@ import android.util.SparseArray
 import android.util.SparseBooleanArray
 import android.view.Display.DEFAULT_DISPLAY
 import android.window.WindowContainerToken
+import androidx.core.util.forEach
+import com.android.internal.protolog.ProtoLog
+import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 
 /** Provides per display window container tokens for [DesktopWallpaperActivity]. */
 class DesktopWallpaperActivityTokenProvider {
@@ -28,6 +31,7 @@ class DesktopWallpaperActivityTokenProvider {
     private val wallpaperActivityVisByDisplayId = SparseBooleanArray()
 
     fun setToken(token: WindowContainerToken, displayId: Int = DEFAULT_DISPLAY) {
+        logV("Setting desktop wallpaper activity token for display %s", displayId)
         wallpaperActivityTokenByDisplayId[displayId] = token
     }
 
@@ -36,9 +40,19 @@ class DesktopWallpaperActivityTokenProvider {
     }
 
     fun removeToken(displayId: Int = DEFAULT_DISPLAY) {
+        logV("Remove desktop wallpaper activity token for display %s", displayId)
         wallpaperActivityTokenByDisplayId.delete(displayId)
     }
 
+    fun removeToken(token: WindowContainerToken) {
+        wallpaperActivityTokenByDisplayId.forEach { displayId, value ->
+            if (value == token) {
+                logV("Remove desktop wallpaper activity token for display %s", displayId)
+                wallpaperActivityTokenByDisplayId.delete(displayId)
+            }
+        }
+    }
+
     fun setWallpaperActivityIsVisible(
         isVisible: Boolean = false,
         displayId: Int = DEFAULT_DISPLAY,
@@ -50,4 +64,12 @@ class DesktopWallpaperActivityTokenProvider {
         return wallpaperActivityTokenByDisplayId[displayId] != null &&
             wallpaperActivityVisByDisplayId.get(displayId, false)
     }
+
+    private fun logV(msg: String, vararg arguments: Any?) {
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
+    companion object {
+        private const val TAG = "DesktopWallpaperActivityTokenProvider"
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
index d2ceb67030fc..ef216b1ae987 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
@@ -30,6 +30,8 @@ import androidx.annotation.NonNull;
 import com.android.internal.protolog.ProtoLog;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.SyncTransactionQueue;
+import com.android.wm.shell.desktopmode.DesktopWallpaperActivity;
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.sysui.ShellInit;
@@ -57,23 +59,30 @@ public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
     private final SyncTransactionQueue mSyncQueue;
     private final Optional<RecentTasksController> mRecentTasksOptional;
     private final Optional<WindowDecorViewModel> mWindowDecorViewModelOptional;
+    private final Optional<DesktopWallpaperActivityTokenProvider>
+            mDesktopWallpaperActivityTokenProviderOptional;
+
     /**
      * This constructor is used by downstream products.
      */
     public FullscreenTaskListener(SyncTransactionQueue syncQueue) {
         this(null /* shellInit */, null /* shellTaskOrganizer */, syncQueue, Optional.empty(),
-                Optional.empty());
+                Optional.empty(), Optional.empty());
     }
 
     public FullscreenTaskListener(ShellInit shellInit,
             ShellTaskOrganizer shellTaskOrganizer,
             SyncTransactionQueue syncQueue,
             Optional<RecentTasksController> recentTasksOptional,
-            Optional<WindowDecorViewModel> windowDecorViewModelOptional) {
+            Optional<WindowDecorViewModel> windowDecorViewModelOptional,
+            Optional<DesktopWallpaperActivityTokenProvider>
+                    desktopWallpaperActivityTokenProviderOptional) {
         mShellTaskOrganizer = shellTaskOrganizer;
         mSyncQueue = syncQueue;
         mRecentTasksOptional = recentTasksOptional;
         mWindowDecorViewModelOptional = windowDecorViewModelOptional;
+        mDesktopWallpaperActivityTokenProviderOptional =
+                desktopWallpaperActivityTokenProviderOptional;
         // Note: Some derivative FullscreenTaskListener implementations do not use ShellInit
         if (shellInit != null) {
             shellInit.addInitCallback(this::onInit, this);
@@ -162,6 +171,12 @@ public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
                 taskInfo.taskId);
         mTasks.remove(taskInfo.taskId);
         mWindowDecorViewModelOptional.ifPresent(v -> v.onTaskVanished(taskInfo));
+        mDesktopWallpaperActivityTokenProviderOptional.ifPresent(
+                provider -> {
+                    if (DesktopWallpaperActivity.isWallpaperTask(taskInfo)) {
+                        provider.removeToken(taskInfo.getToken());
+                    }
+                });
         if (Transitions.ENABLE_SHELL_TRANSITIONS) return;
         if (mWindowDecorViewModelOptional.isPresent()) {
             mWindowDecorViewModelOptional.get().destroyWindowDecoration(taskInfo);