Remove DesktopWallpaperActivity token onVanished.

Sometimes desktop wallpaper activity might be cleared up without a
transition. E.g "clear all" button on overview. When that happens the
token that we store in shell never gets removed as it was only removed
by transition observing.

This cl removes the token when onTaskVanished callback is triggerd as well.

Bug: 391501236
Test: manual
Flag: EXEMPT bugfix
Change-Id: I29be4fd67a2a6a54fc1df035ac31ca0f77483382

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);