[PiP2 on Desktop] Don't exit Desktop Mode when PiP is active.

Recall: http://recall/clips/99739260-beeb-4439-957a-79502a8eb8d1

With PiP on Desktop Mode, we should remain in Desktop Mode in the
following cases:
1) Minimizing the last open Desktop window into PiP
2) While PiP is running, another Desktop window that is the last
   open Desktop window is minimized or closed

In both cases, closing the PiP window should then exit Desktop Mode.
If PiP is running and another Desktop window switches to fullscreen, we
will exit Desktop Mode but PiP will remain, and go to fullscreen when
expanded.

Test: atest DesktopTasksTransitionObserverTest DesktopRepositoryTest
DesktopTasksControllerTest PipSchedulerTest
Flag: com.android.window.flags.enable_desktop_windowing_pip
Bug: 385207728
Bug: 384445396

Change-Id: Ie44842dd2ece08638e4ca2e0f8916e76b8fc4dc6

13 files changed, 595 insertions, 56 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 2600bcc18f72..0495767f3ef1 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
@@ -91,6 +91,7 @@ import com.android.wm.shell.compatui.impl.DefaultComponentIdGenerator;
 import com.android.wm.shell.desktopmode.DesktopMode;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.desktopmode.DesktopUserRepositories;
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider;
 import com.android.wm.shell.displayareahelper.DisplayAreaHelper;
 import com.android.wm.shell.displayareahelper.DisplayAreaHelperController;
 import com.android.wm.shell.freeform.FreeformComponents;
@@ -1033,6 +1034,24 @@ public abstract class WMShellBaseModule {
         });
     }
 
+    @WMSingleton
+    @Provides
+    static DesktopWallpaperActivityTokenProvider provideDesktopWallpaperActivityTokenProvider() {
+        return new DesktopWallpaperActivityTokenProvider();
+    }
+
+    @WMSingleton
+    @Provides
+    static Optional<DesktopWallpaperActivityTokenProvider>
+            provideOptionalDesktopWallpaperActivityTokenProvider(
+            Context context,
+            DesktopWallpaperActivityTokenProvider desktopWallpaperActivityTokenProvider) {
+        if (DesktopModeStatus.canEnterDesktopMode(context)) {
+            return Optional.of(desktopWallpaperActivityTokenProvider);
+        }
+        return Optional.empty();
+    }
+
     //
     // Task Stack
     //
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 1916215dea74..483c8008ba65 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
@@ -1312,12 +1312,6 @@ 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/dagger/pip/Pip2Module.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java
index 9e2b9b20be16..c8d0dab39837 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java
@@ -41,6 +41,7 @@ import com.android.wm.shell.common.pip.SizeSpecSource;
 import com.android.wm.shell.dagger.WMShellBaseModule;
 import com.android.wm.shell.dagger.WMSingleton;
 import com.android.wm.shell.desktopmode.DesktopUserRepositories;
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider;
 import com.android.wm.shell.pip2.phone.PhonePipMenuController;
 import com.android.wm.shell.pip2.phone.PipController;
 import com.android.wm.shell.pip2.phone.PipMotionHelper;
@@ -82,11 +83,14 @@ public abstract class Pip2Module {
             @NonNull PipTransitionState pipStackListenerController,
             @NonNull PipDisplayLayoutState pipDisplayLayoutState,
             @NonNull PipUiStateChangeController pipUiStateChangeController,
-            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional) {
+            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
+            Optional<DesktopWallpaperActivityTokenProvider>
+                    desktopWallpaperActivityTokenProviderOptional) {
         return new PipTransition(context, shellInit, shellTaskOrganizer, transitions,
                 pipBoundsState, null, pipBoundsAlgorithm, pipTaskListener,
                 pipScheduler, pipStackListenerController, pipDisplayLayoutState,
-                pipUiStateChangeController, desktopUserRepositoriesOptional);
+                pipUiStateChangeController, desktopUserRepositoriesOptional,
+                desktopWallpaperActivityTokenProviderOptional);
     }
 
     @WMSingleton
@@ -138,9 +142,12 @@ public abstract class Pip2Module {
             @ShellMainThread ShellExecutor mainExecutor,
             PipTransitionState pipTransitionState,
             Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
+            Optional<DesktopWallpaperActivityTokenProvider>
+                    desktopWallpaperActivityTokenProviderOptional,
             RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
         return new PipScheduler(context, pipBoundsState, mainExecutor, pipTransitionState,
-                desktopUserRepositoriesOptional, rootTaskDisplayAreaOrganizer);
+                desktopUserRepositoriesOptional, desktopWallpaperActivityTokenProviderOptional,
+                rootTaskDisplayAreaOrganizer);
     }
 
     @WMSingleton
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 d3066645f32e..1a58363dab81 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
@@ -27,6 +27,7 @@ import androidx.core.util.forEach
 import androidx.core.util.keyIterator
 import androidx.core.util.valueIterator
 import com.android.internal.protolog.ProtoLog
+import com.android.window.flags.Flags
 import com.android.wm.shell.desktopmode.persistence.DesktopPersistentRepository
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 import com.android.wm.shell.shared.annotations.ShellMainThread
@@ -56,6 +57,10 @@ class DesktopRepository(
      * @property topTransparentFullscreenTaskId the task id of any current top transparent
      *   fullscreen task launched on top of Desktop Mode. Cleared when the transparent task is
      *   closed or sent to back. (top is at index 0).
+     * @property pipTaskId the task id of PiP task entered while in Desktop Mode.
+     * @property pipShouldKeepDesktopActive whether an active PiP window should keep the Desktop
+     *   Mode session active. Only false when we are explicitly exiting Desktop Mode (via user
+     *   action) while there is an active PiP window.
      */
     private data class DesktopTaskData(
         val activeTasks: ArraySet<Int> = ArraySet(),
@@ -66,6 +71,8 @@ class DesktopRepository(
         val freeformTasksInZOrder: ArrayList<Int> = ArrayList(),
         var fullImmersiveTaskId: Int? = null,
         var topTransparentFullscreenTaskId: Int? = null,
+        var pipTaskId: Int? = null,
+        var pipShouldKeepDesktopActive: Boolean = true,
     ) {
         fun deepCopy(): DesktopTaskData =
             DesktopTaskData(
@@ -76,6 +83,8 @@ class DesktopRepository(
                 freeformTasksInZOrder = ArrayList(freeformTasksInZOrder),
                 fullImmersiveTaskId = fullImmersiveTaskId,
                 topTransparentFullscreenTaskId = topTransparentFullscreenTaskId,
+                pipTaskId = pipTaskId,
+                pipShouldKeepDesktopActive = pipShouldKeepDesktopActive,
             )
 
         fun clear() {
@@ -86,6 +95,8 @@ class DesktopRepository(
             freeformTasksInZOrder.clear()
             fullImmersiveTaskId = null
             topTransparentFullscreenTaskId = null
+            pipTaskId = null
+            pipShouldKeepDesktopActive = true
         }
     }
 
@@ -104,6 +115,9 @@ class DesktopRepository(
     /* Tracks last bounds of task before toggled to immersive state. */
     private val boundsBeforeFullImmersiveByTaskId = SparseArray<Rect>()
 
+    /* Callback for when a pending PiP transition has been aborted. */
+    private var onPipAbortedCallback: ((Int, Int) -> Unit)? = null
+
     private var desktopGestureExclusionListener: Consumer<Region>? = null
     private var desktopGestureExclusionExecutor: Executor? = null
 
@@ -302,6 +316,54 @@ class DesktopRepository(
         }
     }
 
+    /** Set whether the given task is the Desktop-entered PiP task in this display. */
+    fun setTaskInPip(displayId: Int, taskId: Int, enterPip: Boolean) {
+        val desktopData = desktopTaskDataByDisplayId.getOrCreate(displayId)
+        if (enterPip) {
+            desktopData.pipTaskId = taskId
+            desktopData.pipShouldKeepDesktopActive = true
+        } else {
+            desktopData.pipTaskId =
+                if (desktopData.pipTaskId == taskId) null
+                else {
+                    logW(
+                        "setTaskInPip: taskId=$taskId did not match saved taskId=${desktopData.pipTaskId}"
+                    )
+                    desktopData.pipTaskId
+                }
+        }
+        notifyVisibleTaskListeners(displayId, getVisibleTaskCount(displayId))
+    }
+
+    /** Returns whether there is a PiP that was entered/minimized from Desktop in this display. */
+    fun isMinimizedPipPresentInDisplay(displayId: Int): Boolean =
+        desktopTaskDataByDisplayId.getOrCreate(displayId).pipTaskId != null
+
+    /** Returns whether the given task is the Desktop-entered PiP task in this display. */
+    fun isTaskMinimizedPipInDisplay(displayId: Int, taskId: Int): Boolean =
+        desktopTaskDataByDisplayId.getOrCreate(displayId).pipTaskId == taskId
+
+    /** Returns whether Desktop session should be active in this display due to active PiP. */
+    fun shouldDesktopBeActiveForPip(displayId: Int): Boolean =
+        Flags.enableDesktopWindowingPip() &&
+            isMinimizedPipPresentInDisplay(displayId) &&
+            desktopTaskDataByDisplayId.getOrCreate(displayId).pipShouldKeepDesktopActive
+
+    /** Saves whether a PiP window should keep Desktop session active in this display. */
+    fun setPipShouldKeepDesktopActive(displayId: Int, keepActive: Boolean) {
+        desktopTaskDataByDisplayId.getOrCreate(displayId).pipShouldKeepDesktopActive = keepActive
+    }
+
+    /** Saves callback to handle a pending PiP transition being aborted. */
+    fun setOnPipAbortedCallback(callbackIfPipAborted: ((Int, Int) -> Unit)?) {
+        onPipAbortedCallback = callbackIfPipAborted
+    }
+
+    /** Invokes callback to handle a pending PiP transition with the given task id being aborted. */
+    fun onPipAborted(displayId: Int, pipTaskId: Int) {
+        onPipAbortedCallback?.invoke(displayId, pipTaskId)
+    }
+
     /** Set whether the given task is the full-immersive task in this display. */
     fun setTaskInFullImmersiveState(displayId: Int, taskId: Int, immersive: Boolean) {
         val desktopData = desktopTaskDataByDisplayId.getOrCreate(displayId)
@@ -338,8 +400,12 @@ class DesktopRepository(
     }
 
     private fun notifyVisibleTaskListeners(displayId: Int, visibleTasksCount: Int) {
+        val visibleAndPipTasksCount =
+            if (shouldDesktopBeActiveForPip(displayId)) visibleTasksCount + 1 else visibleTasksCount
         visibleTasksListeners.forEach { (listener, executor) ->
-            executor.execute { listener.onTasksVisibilityChanged(displayId, visibleTasksCount) }
+            executor.execute {
+                listener.onTasksVisibilityChanged(displayId, visibleAndPipTasksCount)
+            }
         }
     }
 
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 050dfb6f562c..6013648c9806 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
@@ -225,7 +225,6 @@ class DesktopTasksController(
     // Launch cookie used to identify a drag and drop transition to fullscreen after it has begun.
     // Used to prevent handleRequest from moving the new fullscreen task to freeform.
     private var dragAndDropFullscreenCookie: Binder? = null
-    private var pendingPipTransitionAndTask: Pair<IBinder, Int>? = null
 
     init {
         desktopMode = DesktopModeImpl()
@@ -321,18 +320,29 @@ class DesktopTasksController(
     fun visibleTaskCount(displayId: Int): Int = taskRepository.getVisibleTaskCount(displayId)
 
     /**
-     * Returns true if any freeform tasks are visible or if a transparent fullscreen task exists on
-     * top in Desktop Mode.
+     * Returns true if any of the following is true:
+     * - Any freeform tasks are visible
+     * - A transparent fullscreen task exists on top in Desktop Mode
+     * - PiP on Desktop Windowing is enabled, there is an active PiP window and the desktop
+     *   wallpaper is visible.
      */
     fun isDesktopModeShowing(displayId: Int): Boolean {
+        val hasVisibleTasks = visibleTaskCount(displayId) > 0
+        val hasTopTransparentFullscreenTask =
+            taskRepository.getTopTransparentFullscreenTaskId(displayId) != null
+        val hasMinimizedPip =
+            Flags.enableDesktopWindowingPip() &&
+                taskRepository.isMinimizedPipPresentInDisplay(displayId) &&
+                desktopWallpaperActivityTokenProvider.isWallpaperActivityVisible(displayId)
         if (
             DesktopModeFlags.INCLUDE_TOP_TRANSPARENT_FULLSCREEN_TASK_IN_DESKTOP_HEURISTIC
                 .isTrue() && DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_MODALS_POLICY.isTrue()
         ) {
-            return visibleTaskCount(displayId) > 0 ||
-                taskRepository.getTopTransparentFullscreenTaskId(displayId) != null
+            return hasVisibleTasks || hasTopTransparentFullscreenTask || hasMinimizedPip
+        } else if (Flags.enableDesktopWindowingPip()) {
+            return hasVisibleTasks || hasMinimizedPip
         }
-        return visibleTaskCount(displayId) > 0
+        return hasVisibleTasks
     }
 
     /** Moves focused task to desktop mode for given [displayId]. */
@@ -592,7 +602,7 @@ class DesktopTasksController(
     ): ((IBinder) -> Unit)? {
         val taskId = taskInfo.taskId
         desktopTilingDecorViewModel.removeTaskIfTiled(displayId, taskId)
-        performDesktopExitCleanupIfNeeded(taskId, displayId, wct)
+        performDesktopExitCleanupIfNeeded(taskId, displayId, wct, forceToFullscreen = false)
         taskRepository.addClosingTask(displayId, taskId)
         taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(
             doesAnyTaskRequireTaskbarRounding(displayId, taskId)
@@ -624,8 +634,12 @@ class DesktopTasksController(
                 )
             val requestRes = transitions.dispatchRequest(Binder(), requestInfo, /* skip= */ null)
             wct.merge(requestRes.second, true)
-            pendingPipTransitionAndTask =
-                freeformTaskTransitionStarter.startPipTransition(wct) to taskInfo.taskId
+            freeformTaskTransitionStarter.startPipTransition(wct)
+            taskRepository.setTaskInPip(taskInfo.displayId, taskInfo.taskId, enterPip = true)
+            taskRepository.setOnPipAbortedCallback { displayId, taskId ->
+                minimizeTaskInner(shellTaskOrganizer.getRunningTaskInfo(taskId)!!)
+                taskRepository.setTaskInPip(displayId, taskId, enterPip = false)
+            }
             return
         }
 
@@ -636,7 +650,7 @@ class DesktopTasksController(
         val taskId = taskInfo.taskId
         val displayId = taskInfo.displayId
         val wct = WindowContainerTransaction()
-        performDesktopExitCleanupIfNeeded(taskId, displayId, wct)
+        performDesktopExitCleanupIfNeeded(taskId, displayId, wct, forceToFullscreen = false)
         // Notify immersive handler as it might need to exit immersive state.
         val exitResult =
             desktopImmersiveController.exitImmersiveIfApplicable(
@@ -898,7 +912,12 @@ class DesktopTasksController(
         }
 
         if (Flags.enablePerDisplayDesktopWallpaperActivity()) {
-            performDesktopExitCleanupIfNeeded(task.taskId, task.displayId, wct)
+            performDesktopExitCleanupIfNeeded(
+                task.taskId,
+                task.displayId,
+                wct,
+                forceToFullscreen = false,
+            )
         }
 
         transitions.startTransition(TRANSIT_CHANGE, wct, /* handler= */ null)
@@ -1414,7 +1433,9 @@ class DesktopTasksController(
         taskId: Int,
         displayId: Int,
         wct: WindowContainerTransaction,
+        forceToFullscreen: Boolean,
     ) {
+        taskRepository.setPipShouldKeepDesktopActive(displayId, !forceToFullscreen)
         if (Flags.enablePerDisplayDesktopWallpaperActivity()) {
             if (!taskRepository.isOnlyVisibleNonClosingTask(taskId, displayId)) {
                 return
@@ -1422,6 +1443,12 @@ class DesktopTasksController(
             if (displayId != DEFAULT_DISPLAY) {
                 return
             }
+        } else if (
+            Flags.enableDesktopWindowingPip() &&
+                taskRepository.isMinimizedPipPresentInDisplay(displayId) &&
+                !forceToFullscreen
+        ) {
+            return
         } else {
             if (!taskRepository.isOnlyVisibleNonClosingTask(taskId)) {
                 return
@@ -1462,21 +1489,6 @@ class DesktopTasksController(
         return false
     }
 
-    override fun onTransitionConsumed(
-        transition: IBinder,
-        aborted: Boolean,
-        finishT: Transaction?,
-    ) {
-        pendingPipTransitionAndTask?.let { (pipTransition, taskId) ->
-            if (transition == pipTransition) {
-                if (aborted) {
-                    shellTaskOrganizer.getRunningTaskInfo(taskId)?.let { minimizeTaskInner(it) }
-                }
-                pendingPipTransitionAndTask = null
-            }
-        }
-    }
-
     override fun handleRequest(
         transition: IBinder,
         request: TransitionRequestInfo,
@@ -1926,7 +1938,12 @@ class DesktopTasksController(
         if (!isDesktopModeShowing(task.displayId)) return null
 
         val wct = WindowContainerTransaction()
-        performDesktopExitCleanupIfNeeded(task.taskId, task.displayId, wct)
+        performDesktopExitCleanupIfNeeded(
+            task.taskId,
+            task.displayId,
+            wct,
+            forceToFullscreen = false,
+        )
 
         if (!DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION.isTrue()) {
             taskRepository.addClosingTask(task.displayId, task.taskId)
@@ -2053,7 +2070,12 @@ class DesktopTasksController(
             wct.setDensityDpi(taskInfo.token, getDefaultDensityDpi())
         }
 
-        performDesktopExitCleanupIfNeeded(taskInfo.taskId, taskInfo.displayId, wct)
+        performDesktopExitCleanupIfNeeded(
+            taskInfo.taskId,
+            taskInfo.displayId,
+            wct,
+            forceToFullscreen = true,
+        )
     }
 
     private fun cascadeWindow(bounds: Rect, displayLayout: DisplayLayout, displayId: Int) {
@@ -2087,7 +2109,12 @@ class DesktopTasksController(
         // want it overridden in multi-window.
         wct.setDensityDpi(taskInfo.token, getDefaultDensityDpi())
 
-        performDesktopExitCleanupIfNeeded(taskInfo.taskId, taskInfo.displayId, wct)
+        performDesktopExitCleanupIfNeeded(
+            taskInfo.taskId,
+            taskInfo.displayId,
+            wct,
+            forceToFullscreen = false,
+        )
     }
 
     /** Returns the ID of the Task that will be minimized, or null if no task will be minimized. */
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 9bf5555fc194..d61ffdaf5cf8 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
@@ -21,9 +21,11 @@ import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
 import android.content.Context
 import android.os.IBinder
 import android.view.SurfaceControl
-import android.view.WindowManager
 import android.view.WindowManager.TRANSIT_CLOSE
+import android.view.WindowManager.TRANSIT_OPEN
+import android.view.WindowManager.TRANSIT_PIP
 import android.view.WindowManager.TRANSIT_TO_BACK
+import android.view.WindowManager.TRANSIT_TO_FRONT
 import android.window.DesktopModeFlags
 import android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY
 import android.window.TransitionInfo
@@ -39,6 +41,8 @@ import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
+import com.android.wm.shell.transition.Transitions.TRANSIT_EXIT_PIP
+import com.android.wm.shell.transition.Transitions.TRANSIT_REMOVE_PIP
 
 /**
  * A [Transitions.TransitionObserver] that observes shell transitions and updates the
@@ -57,6 +61,8 @@ class DesktopTasksTransitionObserver(
 ) : Transitions.TransitionObserver {
 
     private var transitionToCloseWallpaper: IBinder? = null
+    /* Pending PiP transition and its associated display id and task id. */
+    private var pendingPipTransitionAndPipTask: Triple<IBinder, Int, Int>? = null
     private var currentProfileId: Int
 
     init {
@@ -90,6 +96,33 @@ class DesktopTasksTransitionObserver(
             removeTaskIfNeeded(info)
         }
         removeWallpaperOnLastTaskClosingIfNeeded(transition, info)
+
+        val desktopRepository = desktopUserRepositories.getProfile(currentProfileId)
+        info.changes.forEach { change ->
+            change.taskInfo?.let { taskInfo ->
+                if (
+                    Flags.enableDesktopWindowingPip() &&
+                        desktopRepository.isTaskMinimizedPipInDisplay(
+                            taskInfo.displayId,
+                            taskInfo.taskId,
+                        )
+                ) {
+                    when (info.type) {
+                        TRANSIT_PIP ->
+                            pendingPipTransitionAndPipTask =
+                                Triple(transition, taskInfo.displayId, taskInfo.taskId)
+
+                        TRANSIT_EXIT_PIP,
+                        TRANSIT_REMOVE_PIP ->
+                            desktopRepository.setTaskInPip(
+                                taskInfo.displayId,
+                                taskInfo.taskId,
+                                enterPip = false,
+                            )
+                    }
+                }
+            }
+        }
     }
 
     private fun removeTaskIfNeeded(info: TransitionInfo) {
@@ -252,6 +285,18 @@ class DesktopTasksTransitionObserver(
                 }
             }
             transitionToCloseWallpaper = null
+        } else if (pendingPipTransitionAndPipTask?.first == transition) {
+            val desktopRepository = desktopUserRepositories.getProfile(currentProfileId)
+            if (aborted) {
+                pendingPipTransitionAndPipTask?.let {
+                    desktopRepository.onPipAborted(
+                        /*displayId=*/ it.second,
+                        /* taskId=*/ it.third,
+                    )
+                }
+            }
+            desktopRepository.setOnPipAbortedCallback(null)
+            pendingPipTransitionAndPipTask = null
         }
     }
 
@@ -263,11 +308,15 @@ class DesktopTasksTransitionObserver(
             change.taskInfo?.let { taskInfo ->
                 if (DesktopWallpaperActivity.isWallpaperTask(taskInfo)) {
                     when (change.mode) {
-                        WindowManager.TRANSIT_OPEN -> {
+                        TRANSIT_OPEN -> {
                             desktopWallpaperActivityTokenProvider.setToken(
                                 taskInfo.token,
                                 taskInfo.displayId,
                             )
+                            desktopWallpaperActivityTokenProvider.setWallpaperActivityIsVisible(
+                                isVisible = true,
+                                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.
@@ -278,6 +327,16 @@ class DesktopTasksTransitionObserver(
                         }
                         TRANSIT_CLOSE ->
                             desktopWallpaperActivityTokenProvider.removeToken(taskInfo.displayId)
+                        TRANSIT_TO_FRONT ->
+                            desktopWallpaperActivityTokenProvider.setWallpaperActivityIsVisible(
+                                isVisible = true,
+                                taskInfo.displayId,
+                            )
+                        TRANSIT_TO_BACK ->
+                            desktopWallpaperActivityTokenProvider.setWallpaperActivityIsVisible(
+                                isVisible = false,
+                                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
index a87004c07d43..2bd7a9873a5e 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
@@ -17,6 +17,7 @@
 package com.android.wm.shell.desktopmode.desktopwallpaperactivity
 
 import android.util.SparseArray
+import android.util.SparseBooleanArray
 import android.view.Display.DEFAULT_DISPLAY
 import android.window.WindowContainerToken
 
@@ -24,6 +25,7 @@ import android.window.WindowContainerToken
 class DesktopWallpaperActivityTokenProvider {
 
     private val wallpaperActivityTokenByDisplayId = SparseArray<WindowContainerToken>()
+    private val wallpaperActivityVisByDisplayId = SparseBooleanArray()
 
     fun setToken(token: WindowContainerToken, displayId: Int = DEFAULT_DISPLAY) {
         wallpaperActivityTokenByDisplayId[displayId] = token
@@ -36,4 +38,16 @@ class DesktopWallpaperActivityTokenProvider {
     fun removeToken(displayId: Int = DEFAULT_DISPLAY) {
         wallpaperActivityTokenByDisplayId.delete(displayId)
     }
+
+    fun setWallpaperActivityIsVisible(
+        isVisible: Boolean = false,
+        displayId: Int = DEFAULT_DISPLAY,
+    ) {
+        wallpaperActivityVisByDisplayId.put(displayId, isVisible)
+    }
+
+    fun isWallpaperActivityVisible(displayId: Int = DEFAULT_DISPLAY): Boolean {
+        return wallpaperActivityTokenByDisplayId[displayId] != null &&
+            wallpaperActivityVisByDisplayId.get(displayId, false)
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipScheduler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipScheduler.java
index 7f673d2efc68..ea8dac982703 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipScheduler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipScheduler.java
@@ -40,6 +40,7 @@ import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.pip.PipBoundsState;
 import com.android.wm.shell.desktopmode.DesktopUserRepositories;
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.pip2.PipSurfaceTransactionHelper;
 import com.android.wm.shell.pip2.animation.PipAlphaAnimator;
@@ -59,6 +60,8 @@ public class PipScheduler {
     private final ShellExecutor mMainExecutor;
     private final PipTransitionState mPipTransitionState;
     private final Optional<DesktopUserRepositories> mDesktopUserRepositoriesOptional;
+    private final Optional<DesktopWallpaperActivityTokenProvider>
+            mDesktopWallpaperActivityTokenProviderOptional;
     private final RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
     private PipTransitionController mPipTransitionController;
     private PipSurfaceTransactionHelper.SurfaceControlTransactionFactory
@@ -73,12 +76,16 @@ public class PipScheduler {
             ShellExecutor mainExecutor,
             PipTransitionState pipTransitionState,
             Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
+            Optional<DesktopWallpaperActivityTokenProvider>
+                    desktopWallpaperActivityTokenProviderOptional,
             RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
         mContext = context;
         mPipBoundsState = pipBoundsState;
         mMainExecutor = mainExecutor;
         mPipTransitionState = pipTransitionState;
         mDesktopUserRepositoriesOptional = desktopUserRepositoriesOptional;
+        mDesktopWallpaperActivityTokenProviderOptional =
+                desktopWallpaperActivityTokenProviderOptional;
         mRootTaskDisplayAreaOrganizer = rootTaskDisplayAreaOrganizer;
 
         mSurfaceControlTransactionFactory =
@@ -260,10 +267,18 @@ public class PipScheduler {
 
     /** Returns whether PiP is exiting while we're in desktop mode. */
     private boolean isPipExitingToDesktopMode() {
-        return Flags.enableDesktopWindowingPip() && mDesktopUserRepositoriesOptional.isPresent()
-                && (mDesktopUserRepositoriesOptional.get().getCurrent().getVisibleTaskCount(
-                Objects.requireNonNull(mPipTransitionState.getPipTaskInfo()).displayId) > 0
-                || isDisplayInFreeform());
+        // Early return if PiP in Desktop Windowing is not supported.
+        if (!Flags.enableDesktopWindowingPip() || mDesktopUserRepositoriesOptional.isEmpty()
+                || mDesktopWallpaperActivityTokenProviderOptional.isEmpty()) {
+            return false;
+        }
+        final int displayId = Objects.requireNonNull(
+                mPipTransitionState.getPipTaskInfo()).displayId;
+        return mDesktopUserRepositoriesOptional.get().getCurrent().getVisibleTaskCount(displayId)
+                > 0
+                || mDesktopWallpaperActivityTokenProviderOptional.get().isWallpaperActivityVisible(
+                displayId)
+                || isDisplayInFreeform();
     }
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
index 8061ee9090b6..38015ca6d45f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
@@ -63,7 +63,9 @@ import com.android.wm.shell.common.pip.PipDisplayLayoutState;
 import com.android.wm.shell.common.pip.PipMenuController;
 import com.android.wm.shell.common.pip.PipUtils;
 import com.android.wm.shell.common.split.SplitScreenUtils;
+import com.android.wm.shell.desktopmode.DesktopRepository;
 import com.android.wm.shell.desktopmode.DesktopUserRepositories;
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.pip2.animation.PipAlphaAnimator;
 import com.android.wm.shell.pip2.animation.PipEnterAnimator;
@@ -110,6 +112,8 @@ public class PipTransition extends PipTransitionController implements
     private final PipTransitionState mPipTransitionState;
     private final PipDisplayLayoutState mPipDisplayLayoutState;
     private final Optional<DesktopUserRepositories> mDesktopUserRepositoriesOptional;
+    private final Optional<DesktopWallpaperActivityTokenProvider>
+            mDesktopWallpaperActivityTokenProviderOptional;
 
     //
     // Transition caches
@@ -145,7 +149,9 @@ public class PipTransition extends PipTransitionController implements
             PipTransitionState pipTransitionState,
             PipDisplayLayoutState pipDisplayLayoutState,
             PipUiStateChangeController pipUiStateChangeController,
-            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional) {
+            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
+            Optional<DesktopWallpaperActivityTokenProvider>
+                    desktopWallpaperActivityTokenProviderOptional) {
         super(shellInit, shellTaskOrganizer, transitions, pipBoundsState, pipMenuController,
                 pipBoundsAlgorithm);
 
@@ -157,6 +163,8 @@ public class PipTransition extends PipTransitionController implements
         mPipTransitionState.addPipTransitionStateChangedListener(this);
         mPipDisplayLayoutState = pipDisplayLayoutState;
         mDesktopUserRepositoriesOptional = desktopUserRepositoriesOptional;
+        mDesktopWallpaperActivityTokenProviderOptional =
+                desktopWallpaperActivityTokenProviderOptional;
     }
 
     @Override
@@ -826,13 +834,14 @@ public class PipTransition extends PipTransitionController implements
             return false;
         }
 
-
         // Since opening a new task while in Desktop Mode always first open in Fullscreen
         // until DesktopMode Shell code resolves it to Freeform, PipTransition will get a
         // possibility to handle it also. In this case return false to not have it enter PiP.
         final boolean isInDesktopSession = !mDesktopUserRepositoriesOptional.isEmpty()
-                && mDesktopUserRepositoriesOptional.get().getCurrent().getVisibleTaskCount(
-                pipTask.displayId) > 0;
+                && (mDesktopUserRepositoriesOptional.get().getCurrent().getVisibleTaskCount(
+                        pipTask.displayId) > 0
+                    || mDesktopUserRepositoriesOptional.get().getCurrent()
+                        .isMinimizedPipPresentInDisplay(pipTask.displayId));
         if (isInDesktopSession) {
             return false;
         }
@@ -968,6 +977,27 @@ public class PipTransition extends PipTransitionController implements
                         "Unexpected bundle for " + mPipTransitionState);
                 break;
             case PipTransitionState.EXITED_PIP:
+                final TaskInfo pipTask = mPipTransitionState.getPipTaskInfo();
+                final boolean desktopPipEnabled = Flags.enableDesktopWindowingPip()
+                        && mDesktopUserRepositoriesOptional.isPresent()
+                        && mDesktopWallpaperActivityTokenProviderOptional.isPresent();
+                if (desktopPipEnabled && pipTask != null) {
+                    final DesktopRepository desktopRepository =
+                            mDesktopUserRepositoriesOptional.get().getCurrent();
+                    final boolean wallpaperIsVisible =
+                            mDesktopWallpaperActivityTokenProviderOptional.get()
+                                    .isWallpaperActivityVisible(pipTask.displayId);
+                    if (desktopRepository.getVisibleTaskCount(pipTask.displayId) == 0
+                            && wallpaperIsVisible) {
+                        mTransitions.startTransition(
+                                TRANSIT_TO_BACK,
+                                new WindowContainerTransaction().reorder(
+                                        mDesktopWallpaperActivityTokenProviderOptional.get()
+                                                .getToken(pipTask.displayId), /* onTop= */ false),
+                                null
+                        );
+                    }
+                }
                 mPipTransitionState.setPinnedTaskLeash(null);
                 mPipTransitionState.setPipTaskInfo(null);
                 break;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
index 5629127b8c54..daecccef9344 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
@@ -25,6 +25,7 @@ import android.view.Display.DEFAULT_DISPLAY
 import android.view.Display.INVALID_DISPLAY
 import androidx.test.filters.SmallTest
 import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE
+import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.TestShellExecutor
 import com.android.wm.shell.common.ShellExecutor
@@ -1067,6 +1068,67 @@ class DesktopRepositoryTest : ShellTestCase() {
         assertThat(repo.getTaskInFullImmersiveState(DEFAULT_DESKTOP_ID + 1)).isEqualTo(2)
     }
 
+    @Test
+    fun setTaskInPip_savedAsMinimizedPipInDisplay() {
+        assertThat(repo.isTaskMinimizedPipInDisplay(DEFAULT_DESKTOP_ID, taskId = 1)).isFalse()
+
+        repo.setTaskInPip(DEFAULT_DESKTOP_ID, taskId = 1, enterPip = true)
+
+        assertThat(repo.isTaskMinimizedPipInDisplay(DEFAULT_DESKTOP_ID, taskId = 1)).isTrue()
+    }
+
+    @Test
+    fun removeTaskInPip_removedAsMinimizedPipInDisplay() {
+        repo.setTaskInPip(DEFAULT_DESKTOP_ID, taskId = 1, enterPip = true)
+        assertThat(repo.isTaskMinimizedPipInDisplay(DEFAULT_DESKTOP_ID, taskId = 1)).isTrue()
+
+        repo.setTaskInPip(DEFAULT_DESKTOP_ID, taskId = 1, enterPip = false)
+
+        assertThat(repo.isTaskMinimizedPipInDisplay(DEFAULT_DESKTOP_ID, taskId = 1)).isFalse()
+    }
+
+    @Test
+    fun setTaskInPip_multipleDisplays_bothAreInPip() {
+        repo.setTaskInPip(DEFAULT_DESKTOP_ID, taskId = 1, enterPip = true)
+        repo.setTaskInPip(DEFAULT_DESKTOP_ID + 1, taskId = 2, enterPip = true)
+
+        assertThat(repo.isTaskMinimizedPipInDisplay(DEFAULT_DESKTOP_ID, taskId = 1)).isTrue()
+        assertThat(repo.isTaskMinimizedPipInDisplay(DEFAULT_DESKTOP_ID + 1, taskId = 2)).isTrue()
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun setPipShouldKeepDesktopActive_shouldKeepDesktopActive() {
+        assertThat(repo.shouldDesktopBeActiveForPip(DEFAULT_DESKTOP_ID)).isFalse()
+
+        repo.setTaskInPip(DEFAULT_DESKTOP_ID, taskId = 1, enterPip = true)
+        repo.setPipShouldKeepDesktopActive(DEFAULT_DESKTOP_ID, keepActive = true)
+
+        assertThat(repo.shouldDesktopBeActiveForPip(DEFAULT_DESKTOP_ID)).isTrue()
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun setPipShouldNotKeepDesktopActive_shouldNotKeepDesktopActive() {
+        repo.setTaskInPip(DEFAULT_DESKTOP_ID, taskId = 1, enterPip = true)
+        assertThat(repo.shouldDesktopBeActiveForPip(DEFAULT_DESKTOP_ID)).isTrue()
+
+        repo.setPipShouldKeepDesktopActive(DEFAULT_DESKTOP_ID, keepActive = false)
+
+        assertThat(repo.shouldDesktopBeActiveForPip(DEFAULT_DESKTOP_ID)).isFalse()
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun removeTaskInPip_shouldNotKeepDesktopActive() {
+        repo.setTaskInPip(DEFAULT_DESKTOP_ID, taskId = 1, enterPip = true)
+        assertThat(repo.shouldDesktopBeActiveForPip(DEFAULT_DESKTOP_ID)).isTrue()
+
+        repo.setTaskInPip(DEFAULT_DESKTOP_ID, taskId = 1, enterPip = false)
+
+        assertThat(repo.shouldDesktopBeActiveForPip(DEFAULT_DESKTOP_ID)).isFalse()
+    }
+
     class TestListener : DesktopRepository.ActiveTasksListener {
         var activeChangesOnDefaultDisplay = 0
         var activeChangesOnSecondaryDisplay = 0
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 692b50303038..4bb743079861 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
@@ -82,6 +82,7 @@ import com.android.dx.mockito.inline.extended.StaticMockitoSession
 import com.android.internal.jank.InteractionJankMonitor
 import com.android.window.flags.Flags
 import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE
+import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP
 import com.android.window.flags.Flags.FLAG_ENABLE_DISPLAY_FOCUS_IN_SHELL_TRANSITIONS
 import com.android.window.flags.Flags.FLAG_ENABLE_FULLY_IMMERSIVE_IN_DESKTOP
 import com.android.window.flags.Flags.FLAG_ENABLE_MOVE_TO_NEXT_DISPLAY_SHORTCUT
@@ -569,6 +570,38 @@ class DesktopTasksControllerTest : ShellTestCase() {
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun isDesktopModeShowing_minimizedPipTask_wallpaperVisible_returnsTrue() {
+        val pipTask = setUpPipTask(autoEnterEnabled = true)
+        whenever(desktopWallpaperActivityTokenProvider.isWallpaperActivityVisible())
+            .thenReturn(true)
+
+        taskRepository.setTaskInPip(DEFAULT_DISPLAY, pipTask.taskId, enterPip = true)
+
+        assertThat(controller.isDesktopModeShowing(displayId = DEFAULT_DISPLAY)).isTrue()
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun isDesktopModeShowing_minimizedPipTask_wallpaperNotVisible_returnsFalse() {
+        val pipTask = setUpPipTask(autoEnterEnabled = true)
+        whenever(desktopWallpaperActivityTokenProvider.isWallpaperActivityVisible())
+            .thenReturn(false)
+
+        taskRepository.setTaskInPip(DEFAULT_DISPLAY, pipTask.taskId, enterPip = true)
+
+        assertThat(controller.isDesktopModeShowing(displayId = DEFAULT_DISPLAY)).isFalse()
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun isDesktopModeShowing_pipTaskNotMinimizedNorVisible_returnsFalse() {
+        setUpPipTask(autoEnterEnabled = true)
+
+        assertThat(controller.isDesktopModeShowing(displayId = DEFAULT_DISPLAY)).isFalse()
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun showDesktopApps_onSecondaryDisplay_desktopWallpaperEnabled_shouldNotShowWallpaper() {
         val homeTask = setUpHomeTask(SECOND_DISPLAY)
@@ -2039,6 +2072,41 @@ class DesktopTasksControllerTest : ShellTestCase() {
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun onDesktopWindowClose_minimizedPipPresent_doesNotExitDesktop() {
+        val freeformTask = setUpFreeformTask().apply { isFocused = true }
+        val pipTask = setUpPipTask(autoEnterEnabled = true)
+
+        taskRepository.setTaskInPip(DEFAULT_DISPLAY, pipTask.taskId, enterPip = true)
+        val wct = WindowContainerTransaction()
+        controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, freeformTask)
+
+        verifyExitDesktopWCTNotExecuted()
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun onDesktopWindowClose_minimizedPipNotPresent_exitDesktop() {
+        val freeformTask = setUpFreeformTask()
+        val pipTask = setUpPipTask(autoEnterEnabled = true)
+        val handler = mock(TransitionHandler::class.java)
+        whenever(transitions.dispatchRequest(any(), any(), anyOrNull()))
+            .thenReturn(android.util.Pair(handler, WindowContainerTransaction()))
+
+        controller.minimizeTask(pipTask)
+        verifyExitDesktopWCTNotExecuted()
+
+        taskRepository.setTaskInPip(DEFAULT_DISPLAY, pipTask.taskId, enterPip = false)
+        val wct = WindowContainerTransaction()
+        controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, freeformTask)
+
+        // Remove wallpaper operation
+        wct.hierarchyOps.any { hop ->
+            hop.type == HIERARCHY_OP_TYPE_REMOVE_TASK && hop.container == wallpaperToken.asBinder()
+        }
+    }
+
+    @Test
     fun onDesktopWindowMinimize_noActiveTask_doesntRemoveWallpaper() {
         val task = setUpFreeformTask(active = false)
         val transition = Binder()
@@ -2055,10 +2123,9 @@ class DesktopTasksControllerTest : ShellTestCase() {
     }
 
     @Test
-    fun onDesktopWindowMinimize_pipTask_autoEnterEnabled_startPipTransition() {
+    fun onPipTaskMinimize_autoEnterEnabled_startPipTransition() {
         val task = setUpPipTask(autoEnterEnabled = true)
         val handler = mock(TransitionHandler::class.java)
-        whenever(freeformTaskTransitionStarter.startPipTransition(any())).thenReturn(Binder())
         whenever(transitions.dispatchRequest(any(), any(), anyOrNull()))
             .thenReturn(android.util.Pair(handler, WindowContainerTransaction()))
 
@@ -2069,7 +2136,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     }
 
     @Test
-    fun onDesktopWindowMinimize_pipTask_autoEnterDisabled_startMinimizeTransition() {
+    fun onPipTaskMinimize_autoEnterDisabled_startMinimizeTransition() {
         val task = setUpPipTask(autoEnterEnabled = false)
         whenever(freeformTaskTransitionStarter.startMinimizedModeTransition(any()))
             .thenReturn(Binder())
@@ -2081,6 +2148,22 @@ class DesktopTasksControllerTest : ShellTestCase() {
     }
 
     @Test
+    fun onPipTaskMinimize_doesntRemoveWallpaper() {
+        val task = setUpPipTask(autoEnterEnabled = true)
+        val handler = mock(TransitionHandler::class.java)
+        whenever(transitions.dispatchRequest(any(), any(), anyOrNull()))
+            .thenReturn(android.util.Pair(handler, WindowContainerTransaction()))
+
+        controller.minimizeTask(task)
+
+        val captor = ArgumentCaptor.forClass(WindowContainerTransaction::class.java)
+        verify(freeformTaskTransitionStarter).startPipTransition(captor.capture())
+        captor.value.hierarchyOps.none { hop ->
+            hop.type == HIERARCHY_OP_TYPE_REMOVE_TASK && hop.container == wallpaperToken.asBinder()
+        }
+    }
+
+    @Test
     fun onDesktopWindowMinimize_singleActiveTask_noWallpaperActivityToken_doesntRemoveWallpaper() {
         val task = setUpFreeformTask(active = true)
         val transition = Binder()
@@ -3125,6 +3208,31 @@ class DesktopTasksControllerTest : ShellTestCase() {
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun moveFocusedTaskToFullscreen_minimizedPipPresent_removeWallpaperActivity() {
+        val freeformTask = setUpFreeformTask()
+        val pipTask = setUpPipTask(autoEnterEnabled = true)
+        val handler = mock(TransitionHandler::class.java)
+        whenever(transitions.dispatchRequest(any(), any(), anyOrNull()))
+            .thenReturn(android.util.Pair(handler, WindowContainerTransaction()))
+
+        controller.minimizeTask(pipTask)
+        verifyExitDesktopWCTNotExecuted()
+
+        freeformTask.isFocused = true
+        controller.enterFullscreen(DEFAULT_DISPLAY, transitionSource = UNKNOWN)
+
+        val wct = getLatestExitDesktopWct()
+        val taskChange = assertNotNull(wct.changes[freeformTask.token.asBinder()])
+        assertThat(taskChange.windowingMode)
+            .isEqualTo(WINDOWING_MODE_UNDEFINED) // inherited FULLSCREEN
+        // Remove wallpaper operation
+        wct.hierarchyOps.any { hop ->
+            hop.type == HIERARCHY_OP_TYPE_REMOVE_TASK && hop.container == wallpaperToken.asBinder()
+        }
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
     fun removeDesktop_multipleTasks_removesAll() {
         val task1 = setUpFreeformTask()
@@ -4851,7 +4959,8 @@ class DesktopTasksControllerTest : ShellTestCase() {
     }
 
     private fun setUpPipTask(autoEnterEnabled: Boolean): RunningTaskInfo {
-        return setUpFreeformTask().apply {
+        // active = false marks the task as non-visible; PiP window doesn't count as visible tasks
+        return setUpFreeformTask(active = false).apply {
             pictureInPictureParams =
                 PictureInPictureParams.Builder().setAutoEnterEnabled(autoEnterEnabled).build()
         }
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 89ab65a42bbf..96ed214e7f88 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
@@ -22,6 +22,7 @@ import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
 import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
+import android.os.Binder
 import android.os.IBinder
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
@@ -29,7 +30,9 @@ import android.view.Display.DEFAULT_DISPLAY
 import android.view.WindowManager
 import android.view.WindowManager.TRANSIT_CLOSE
 import android.view.WindowManager.TRANSIT_OPEN
+import android.view.WindowManager.TRANSIT_PIP
 import android.view.WindowManager.TRANSIT_TO_BACK
+import android.view.WindowManager.TRANSIT_TO_FRONT
 import android.window.IWindowContainerToken
 import android.window.TransitionInfo
 import android.window.TransitionInfo.Change
@@ -38,6 +41,7 @@ import android.window.WindowContainerTransaction
 import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER
 import com.android.modules.utils.testing.ExtendedMockitoRule
 import com.android.window.flags.Flags
+import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP
 import com.android.wm.shell.MockToken
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.back.BackAnimationController
@@ -47,6 +51,8 @@ import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpape
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
+import com.android.wm.shell.transition.Transitions.TRANSIT_EXIT_PIP
+import com.android.wm.shell.transition.Transitions.TRANSIT_REMOVE_PIP
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import org.junit.Before
@@ -300,6 +306,115 @@ class DesktopTasksTransitionObserverTest {
         verify(taskRepository).clearTopTransparentFullscreenTaskId(topTransparentTask.displayId)
     }
 
+    @Test
+    fun transitOpenWallpaper_wallpaperActivityVisibilitySaved() {
+        val wallpaperTask = createWallpaperTaskInfo()
+
+        transitionObserver.onTransitionReady(
+            transition = mock(),
+            info = createOpenChangeTransition(wallpaperTask),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+        )
+
+        verify(desktopWallpaperActivityTokenProvider)
+            .setWallpaperActivityIsVisible(isVisible = true, wallpaperTask.displayId)
+    }
+
+    @Test
+    fun transitToFrontWallpaper_wallpaperActivityVisibilitySaved() {
+        val wallpaperTask = createWallpaperTaskInfo()
+
+        transitionObserver.onTransitionReady(
+            transition = mock(),
+            info = createToFrontTransition(wallpaperTask),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+        )
+
+        verify(desktopWallpaperActivityTokenProvider)
+            .setWallpaperActivityIsVisible(isVisible = true, wallpaperTask.displayId)
+    }
+
+    @Test
+    fun transitToBackWallpaper_wallpaperActivityVisibilitySaved() {
+        val wallpaperTask = createWallpaperTaskInfo()
+
+        transitionObserver.onTransitionReady(
+            transition = mock(),
+            info = createToBackTransition(wallpaperTask),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+        )
+
+        verify(desktopWallpaperActivityTokenProvider)
+            .setWallpaperActivityIsVisible(isVisible = false, wallpaperTask.displayId)
+    }
+
+    @Test
+    fun transitCloseWallpaper_wallpaperActivityVisibilitySaved() {
+        val wallpaperTask = createWallpaperTaskInfo()
+
+        transitionObserver.onTransitionReady(
+            transition = mock(),
+            info = createCloseTransition(wallpaperTask),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+        )
+
+        verify(desktopWallpaperActivityTokenProvider).removeToken(wallpaperTask.displayId)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun pendingPipTransitionAborted_taskRepositoryOnPipAbortedInvoked() {
+        val task = createTaskInfo(1, WINDOWING_MODE_FREEFORM)
+        val pipTransition = Binder()
+        whenever(taskRepository.isTaskMinimizedPipInDisplay(any(), any())).thenReturn(true)
+
+        transitionObserver.onTransitionReady(
+            transition = pipTransition,
+            info = createOpenChangeTransition(task, TRANSIT_PIP),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+        )
+        transitionObserver.onTransitionFinished(transition = pipTransition, aborted = true)
+
+        verify(taskRepository).onPipAborted(task.displayId, task.taskId)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun exitPipTransition_taskRepositoryClearTaskInPip() {
+        val task = createTaskInfo(1, WINDOWING_MODE_FREEFORM)
+        whenever(taskRepository.isTaskMinimizedPipInDisplay(any(), any())).thenReturn(true)
+
+        transitionObserver.onTransitionReady(
+            transition = mock(),
+            info = createOpenChangeTransition(task, type = TRANSIT_EXIT_PIP),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+        )
+
+        verify(taskRepository).setTaskInPip(task.displayId, task.taskId, enterPip = false)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun removePipTransition_taskRepositoryClearTaskInPip() {
+        val task = createTaskInfo(1, WINDOWING_MODE_FREEFORM)
+        whenever(taskRepository.isTaskMinimizedPipInDisplay(any(), any())).thenReturn(true)
+
+        transitionObserver.onTransitionReady(
+            transition = mock(),
+            info = createOpenChangeTransition(task, type = TRANSIT_REMOVE_PIP),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+        )
+
+        verify(taskRepository).setTaskInPip(task.displayId, task.taskId, enterPip = false)
+    }
+
     private fun createBackNavigationTransition(
         task: RunningTaskInfo?,
         type: Int = TRANSIT_TO_BACK,
@@ -331,7 +446,7 @@ class DesktopTasksTransitionObserverTest {
         task: RunningTaskInfo?,
         type: Int = TRANSIT_OPEN,
     ): TransitionInfo {
-        return TransitionInfo(TRANSIT_OPEN, /* flags= */ 0).apply {
+        return TransitionInfo(type, /* flags= */ 0).apply {
             addChange(
                 Change(mock(), mock()).apply {
                     mode = TRANSIT_OPEN
@@ -369,6 +484,19 @@ class DesktopTasksTransitionObserverTest {
         }
     }
 
+    private fun createToFrontTransition(task: RunningTaskInfo?): TransitionInfo {
+        return TransitionInfo(TRANSIT_TO_FRONT, 0 /* flags */).apply {
+            addChange(
+                Change(mock(), mock()).apply {
+                    mode = TRANSIT_TO_FRONT
+                    parent = null
+                    taskInfo = task
+                    flags = flags
+                }
+            )
+        }
+    }
+
     private fun getLatestWct(
         @WindowManager.TransitionType type: Int = TRANSIT_OPEN,
         handlerClass: Class<out Transitions.TransitionHandler>? = null,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/PipSchedulerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/PipSchedulerTest.java
index a8aa25700c7e..c42f6c35bcb0 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/PipSchedulerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/PipSchedulerTest.java
@@ -30,6 +30,7 @@ import static org.mockito.kotlin.MatchersKt.eq;
 import static org.mockito.kotlin.VerificationKt.times;
 import static org.mockito.kotlin.VerificationKt.verify;
 
+import android.app.TaskInfo;
 import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Matrix;
@@ -45,7 +46,9 @@ import androidx.test.filters.SmallTest;
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.pip.PipBoundsState;
+import com.android.wm.shell.desktopmode.DesktopRepository;
 import com.android.wm.shell.desktopmode.DesktopUserRepositories;
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.pip2.PipSurfaceTransactionHelper;
 import com.android.wm.shell.pip2.animation.PipAlphaAnimator;
@@ -56,6 +59,7 @@ import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 import java.util.Optional;
@@ -83,7 +87,8 @@ public class PipSchedulerTest {
     @Mock private PipSurfaceTransactionHelper.SurfaceControlTransactionFactory mMockFactory;
     @Mock private SurfaceControl.Transaction mMockTransaction;
     @Mock private PipAlphaAnimator mMockAlphaAnimator;
-    @Mock private Optional<DesktopUserRepositories> mMockOptionalDesktopUserRepositories;
+    @Mock private DesktopUserRepositories mMockDesktopUserRepositories;
+    @Mock private DesktopWallpaperActivityTokenProvider mMockDesktopWallpaperActivityTokenProvider;
     @Mock private RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
 
     @Captor private ArgumentCaptor<Runnable> mRunnableArgumentCaptor;
@@ -100,9 +105,13 @@ public class PipSchedulerTest {
         when(mMockFactory.getTransaction()).thenReturn(mMockTransaction);
         when(mMockTransaction.setMatrix(any(SurfaceControl.class), any(Matrix.class), any()))
                 .thenReturn(mMockTransaction);
+        when(mMockDesktopUserRepositories.getCurrent())
+                .thenReturn(Mockito.mock(DesktopRepository.class));
+        when(mMockPipTransitionState.getPipTaskInfo()).thenReturn(Mockito.mock(TaskInfo.class));
 
         mPipScheduler = new PipScheduler(mMockContext, mMockPipBoundsState, mMockMainExecutor,
-                mMockPipTransitionState, mMockOptionalDesktopUserRepositories,
+                mMockPipTransitionState, Optional.of(mMockDesktopUserRepositories),
+                Optional.of(mMockDesktopWallpaperActivityTokenProvider),
                 mRootTaskDisplayAreaOrganizer);
         mPipScheduler.setPipTransitionController(mMockPipTransitionController);
         mPipScheduler.setSurfaceControlTransactionFactory(mMockFactory);