Merge changes from topic "desktop-pip-shellmodule" into main

* changes:
  Fix DesktopPipTransitionObserver dependencies in WMShellModule.
  Revert^2 "[PiP on Desktop] Minimizing last task to PiP should exit Desktop"

9 files changed, 399 insertions, 252 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 257e27a338be..67a4d6cf89bf 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
@@ -98,6 +98,7 @@ import com.android.wm.shell.desktopmode.DesktopModeKeyGestureHandler;
 import com.android.wm.shell.desktopmode.DesktopModeLoggerTransitionObserver;
 import com.android.wm.shell.desktopmode.DesktopModeMoveToDisplayTransitionHandler;
 import com.android.wm.shell.desktopmode.DesktopModeUiEventLogger;
+import com.android.wm.shell.desktopmode.DesktopPipTransitionObserver;
 import com.android.wm.shell.desktopmode.DesktopTaskChangeListener;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.desktopmode.DesktopTasksLimiter;
@@ -780,6 +781,7 @@ public abstract class WMShellModule {
             OverviewToDesktopTransitionObserver overviewToDesktopTransitionObserver,
             DesksOrganizer desksOrganizer,
             Optional<DesksTransitionObserver> desksTransitionObserver,
+            Optional<DesktopPipTransitionObserver> desktopPipTransitionObserver,
             UserProfileContexts userProfileContexts,
             DesktopModeCompatPolicy desktopModeCompatPolicy,
             DragToDisplayTransitionHandler dragToDisplayTransitionHandler,
@@ -823,6 +825,7 @@ public abstract class WMShellModule {
                 overviewToDesktopTransitionObserver,
                 desksOrganizer,
                 desksTransitionObserver.get(),
+                desktopPipTransitionObserver,
                 userProfileContexts,
                 desktopModeCompatPolicy,
                 dragToDisplayTransitionHandler,
@@ -1225,6 +1228,7 @@ public abstract class WMShellModule {
             Transitions transitions,
             ShellTaskOrganizer shellTaskOrganizer,
             Optional<DesktopMixedTransitionHandler> desktopMixedTransitionHandler,
+            Optional<DesktopPipTransitionObserver> desktopPipTransitionObserver,
             Optional<BackAnimationController> backAnimationController,
             DesktopWallpaperActivityTokenProvider desktopWallpaperActivityTokenProvider,
             ShellInit shellInit) {
@@ -1237,6 +1241,7 @@ public abstract class WMShellModule {
                                         transitions,
                                         shellTaskOrganizer,
                                         desktopMixedTransitionHandler.get(),
+                                        desktopPipTransitionObserver,
                                         backAnimationController.get(),
                                         desktopWallpaperActivityTokenProvider,
                                         shellInit)));
@@ -1258,6 +1263,19 @@ public abstract class WMShellModule {
 
     @WMSingleton
     @Provides
+    static Optional<DesktopPipTransitionObserver> provideDesktopPipTransitionObserver(
+            Context context
+    ) {
+        if (DesktopModeStatus.canEnterDesktopMode(context)
+                && DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PIP.isTrue()) {
+            return Optional.of(
+                    new DesktopPipTransitionObserver());
+        }
+        return Optional.empty();
+    }
+
+    @WMSingleton
+    @Provides
     static Optional<DesktopMixedTransitionHandler> provideDesktopMixedTransitionHandler(
             Context context,
             Transitions transitions,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopPipTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopPipTransitionObserver.kt
new file mode 100644
index 000000000000..efd3866e1bc4
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopPipTransitionObserver.kt
@@ -0,0 +1,81 @@
+/*
+ * Copyright (C) 2025 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
+
+import android.app.WindowConfiguration.WINDOWING_MODE_PINNED
+import android.os.IBinder
+import android.window.DesktopModeFlags
+import android.window.TransitionInfo
+import com.android.internal.protolog.ProtoLog
+import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
+
+/**
+ * Observer of PiP in Desktop Mode transitions. At the moment, this is specifically tracking a PiP
+ * transition for a task that is entering PiP via the minimize button on the caption bar.
+ */
+class DesktopPipTransitionObserver {
+    private val pendingPipTransitions = mutableMapOf<IBinder, PendingPipTransition>()
+
+    /** Adds a pending PiP transition to be tracked. */
+    fun addPendingPipTransition(transition: PendingPipTransition) {
+        if (!DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PIP.isTrue) return
+        pendingPipTransitions[transition.token] = transition
+    }
+
+    /**
+     * Called when any transition is ready, which may include transitions not tracked by this
+     * observer.
+     */
+    fun onTransitionReady(transition: IBinder, info: TransitionInfo) {
+        if (!DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PIP.isTrue) return
+        val pipTransition = pendingPipTransitions.remove(transition) ?: return
+
+        logD("Desktop PiP transition ready: %s", transition)
+        for (change in info.changes) {
+            val taskInfo = change.taskInfo
+            if (taskInfo == null || taskInfo.taskId == -1) {
+                continue
+            }
+
+            if (
+                taskInfo.taskId == pipTransition.taskId &&
+                    taskInfo.windowingMode == WINDOWING_MODE_PINNED
+            ) {
+                logD("Desktop PiP transition was successful")
+                pipTransition.onSuccess()
+                return
+            }
+        }
+        logD("Change with PiP task not found in Desktop PiP transition; likely failed")
+    }
+
+    /**
+     * Data tracked for a pending PiP transition.
+     *
+     * @property token the PiP transition that is started.
+     * @property taskId task id of the task entering PiP.
+     * @property onSuccess callback to be invoked if the PiP transition is successful.
+     */
+    data class PendingPipTransition(val token: IBinder, val taskId: Int, val onSuccess: () -> Unit)
+
+    private fun logD(msg: String, vararg arguments: Any?) {
+        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
+    private companion object {
+        private const val TAG = "DesktopPipTransitionObserver"
+    }
+}
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 a4e9c52ac9d9..6cb26b54e802 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
@@ -68,7 +68,6 @@ class DesktopRepository(
      * @property topTransparentFullscreenTaskId the task id of any current top transparent
      *   fullscreen task launched on top of the desk. 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.
      */
     private data class Desk(
         val deskId: Int,
@@ -81,7 +80,6 @@ class DesktopRepository(
         val freeformTasksInZOrder: ArrayList<Int> = ArrayList(),
         var fullImmersiveTaskId: Int? = null,
         var topTransparentFullscreenTaskId: Int? = null,
-        var pipTaskId: Int? = null,
     ) {
         fun deepCopy(): Desk =
             Desk(
@@ -94,7 +92,6 @@ class DesktopRepository(
                 freeformTasksInZOrder = ArrayList(freeformTasksInZOrder),
                 fullImmersiveTaskId = fullImmersiveTaskId,
                 topTransparentFullscreenTaskId = topTransparentFullscreenTaskId,
-                pipTaskId = pipTaskId,
             )
 
         // TODO: b/362720497 - remove when multi-desktops is enabled where instances aren't
@@ -107,7 +104,6 @@ class DesktopRepository(
             freeformTasksInZOrder.clear()
             fullImmersiveTaskId = null
             topTransparentFullscreenTaskId = null
-            pipTaskId = null
         }
     }
 
@@ -127,9 +123,6 @@ 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
 
@@ -611,57 +604,6 @@ class DesktopRepository(
     }
 
     /**
-     * Set whether the given task is the Desktop-entered PiP task in this display's active desk.
-     *
-     * TODO: b/389960283 - add explicit [deskId] argument.
-     */
-    fun setTaskInPip(displayId: Int, taskId: Int, enterPip: Boolean) {
-        val activeDesk =
-            desktopData.getActiveDesk(displayId)
-                ?: error("Expected active desk in display: $displayId")
-        if (enterPip) {
-            activeDesk.pipTaskId = taskId
-        } else {
-            activeDesk.pipTaskId =
-                if (activeDesk.pipTaskId == taskId) null
-                else {
-                    logW(
-                        "setTaskInPip: taskId=%d did not match saved taskId=%d",
-                        taskId,
-                        activeDesk.pipTaskId,
-                    )
-                    activeDesk.pipTaskId
-                }
-        }
-    }
-
-    /**
-     * Returns whether the given task is the Desktop-entered PiP task in this display's active desk.
-     *
-     * TODO: b/389960283 - add explicit [deskId] argument.
-     */
-    fun isTaskMinimizedPipInDisplay(displayId: Int, taskId: Int): Boolean =
-        desktopData.getActiveDesk(displayId)?.pipTaskId == taskId
-
-    /**
-     * Saves callback to handle a pending PiP transition being aborted.
-     *
-     * TODO: b/389960283 - add explicit [deskId] argument.
-     */
-    fun setOnPipAbortedCallback(callbackIfPipAborted: ((displayId: Int, pipTaskId: Int) -> Unit)?) {
-        onPipAbortedCallback = callbackIfPipAborted
-    }
-
-    /**
-     * Invokes callback to handle a pending PiP transition with the given task id being aborted.
-     *
-     * TODO: b/389960283 - add explicit [deskId] argument.
-     */
-    fun onPipAborted(displayId: Int, pipTaskId: Int) {
-        onPipAbortedCallback?.invoke(displayId, pipTaskId)
-    }
-
-    /**
      * Set whether the given task is the full-immersive task in this display's active desk.
      *
      * TODO: b/389960283 - consider forcing callers to use [setTaskInFullImmersiveStateInDesk] with
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 5e9cd9016d92..50f5beb4166a 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
@@ -214,6 +214,7 @@ class DesktopTasksController(
     private val overviewToDesktopTransitionObserver: OverviewToDesktopTransitionObserver,
     private val desksOrganizer: DesksOrganizer,
     private val desksTransitionObserver: DesksTransitionObserver,
+    private val desktopPipTransitionObserver: Optional<DesktopPipTransitionObserver>,
     private val userProfileContexts: UserProfileContexts,
     private val desktopModeCompatPolicy: DesktopModeCompatPolicy,
     private val dragToDisplayTransitionHandler: DragToDisplayTransitionHandler,
@@ -793,10 +794,31 @@ class DesktopTasksController(
 
     fun minimizeTask(taskInfo: RunningTaskInfo, minimizeReason: MinimizeReason) {
         val wct = WindowContainerTransaction()
-
+        val taskId = taskInfo.taskId
+        val displayId = taskInfo.displayId
+        val deskId =
+            taskRepository.getDeskIdForTask(taskInfo.taskId)
+                ?: if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
+                    logW("minimizeTask: desk not found for task: ${taskInfo.taskId}")
+                    return
+                } else {
+                    getDefaultDeskId(taskInfo.displayId)
+                }
+        val isLastTask =
+            if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
+                taskRepository.isOnlyVisibleNonClosingTaskInDesk(
+                    taskId = taskId,
+                    deskId = checkNotNull(deskId) { "Expected non-null deskId" },
+                    displayId = displayId,
+                )
+            } else {
+                taskRepository.isOnlyVisibleNonClosingTask(taskId = taskId, displayId = displayId)
+            }
         val isMinimizingToPip =
             DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PIP.isTrue &&
-                (taskInfo.pictureInPictureParams?.isAutoEnterEnabled() ?: false)
+                desktopPipTransitionObserver.isPresent &&
+                (taskInfo.pictureInPictureParams?.isAutoEnterEnabled ?: false)
+
         // If task is going to PiP, start a PiP transition instead of a minimize transition
         if (isMinimizingToPip) {
             val requestInfo =
@@ -810,75 +832,60 @@ class DesktopTasksController(
                 )
             val requestRes = transitions.dispatchRequest(Binder(), requestInfo, /* skip= */ null)
             wct.merge(requestRes.second, true)
-            freeformTaskTransitionStarter.startPipTransition(wct)
-            taskRepository.setTaskInPip(taskInfo.displayId, taskInfo.taskId, enterPip = true)
-            taskRepository.setOnPipAbortedCallback { displayId, taskId ->
-                minimizeTaskInner(shellTaskOrganizer.getRunningTaskInfo(taskId)!!, minimizeReason)
-                taskRepository.setTaskInPip(displayId, taskId, enterPip = false)
-            }
-            return
-        }
-
-        minimizeTaskInner(taskInfo, minimizeReason)
-    }
 
-    private fun minimizeTaskInner(taskInfo: RunningTaskInfo, minimizeReason: MinimizeReason) {
-        val taskId = taskInfo.taskId
-        val deskId = taskRepository.getDeskIdForTask(taskInfo.taskId)
-        if (deskId == null && DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-            logW("minimizeTaskInner: desk not found for task: ${taskInfo.taskId}")
-            return
-        }
-        val displayId = taskInfo.displayId
-        val wct = WindowContainerTransaction()
-
-        snapEventHandler.removeTaskIfTiled(displayId, taskId)
-        val willExitDesktop = willExitDesktop(taskId, displayId, forceExitDesktop = false)
-        val desktopExitRunnable =
-            performDesktopExitCleanUp(
-                wct = wct,
-                deskId = deskId,
-                displayId = displayId,
-                willExitDesktop = willExitDesktop,
-            )
-        // Notify immersive handler as it might need to exit immersive state.
-        val exitResult =
-            desktopImmersiveController.exitImmersiveIfApplicable(
-                wct = wct,
-                taskInfo = taskInfo,
-                reason = DesktopImmersiveController.ExitReason.MINIMIZED,
-            )
-        if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-            desksOrganizer.minimizeTask(
-                wct = wct,
-                deskId = checkNotNull(deskId) { "Expected non-null deskId" },
-                task = taskInfo,
+            desktopPipTransitionObserver.get().addPendingPipTransition(
+                DesktopPipTransitionObserver.PendingPipTransition(
+                    token = freeformTaskTransitionStarter.startPipTransition(wct),
+                    taskId = taskInfo.taskId,
+                    onSuccess = {
+                        onDesktopTaskEnteredPip(
+                            taskId = taskId,
+                            deskId = deskId,
+                            displayId = taskInfo.displayId,
+                            taskIsLastVisibleTaskBeforePip = isLastTask,
+                        )
+                    },
+                )
             )
         } else {
-            wct.reorder(taskInfo.token, /* onTop= */ false)
-        }
-        val isLastTask =
+            snapEventHandler.removeTaskIfTiled(displayId, taskId)
+            val willExitDesktop = willExitDesktop(taskId, displayId, forceExitDesktop = false)
+            val desktopExitRunnable =
+                performDesktopExitCleanUp(
+                    wct = wct,
+                    deskId = deskId,
+                    displayId = displayId,
+                    willExitDesktop = willExitDesktop,
+                )
+            // Notify immersive handler as it might need to exit immersive state.
+            val exitResult =
+                desktopImmersiveController.exitImmersiveIfApplicable(
+                    wct = wct,
+                    taskInfo = taskInfo,
+                    reason = DesktopImmersiveController.ExitReason.MINIMIZED,
+                )
             if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-                taskRepository.isOnlyVisibleNonClosingTaskInDesk(
-                    taskId = taskId,
+                desksOrganizer.minimizeTask(
+                    wct = wct,
                     deskId = checkNotNull(deskId) { "Expected non-null deskId" },
-                    displayId = displayId,
+                    task = taskInfo,
                 )
             } else {
-                taskRepository.isOnlyVisibleNonClosingTask(taskId = taskId, displayId = displayId)
+                wct.reorder(taskInfo.token, /* onTop= */ false)
             }
-        val transition =
-            freeformTaskTransitionStarter.startMinimizedModeTransition(wct, taskId, isLastTask)
-        desktopTasksLimiter.ifPresent {
-            it.addPendingMinimizeChange(
-                transition = transition,
-                displayId = displayId,
-                taskId = taskId,
-                minimizeReason = minimizeReason,
-            )
+            val transition =
+                freeformTaskTransitionStarter.startMinimizedModeTransition(wct, taskId, isLastTask)
+            desktopTasksLimiter.ifPresent {
+                it.addPendingMinimizeChange(
+                    transition = transition,
+                    displayId = displayId,
+                    taskId = taskId,
+                    minimizeReason = minimizeReason,
+                )
+            }
+            exitResult.asExit()?.runOnTransitionStart?.invoke(transition)
+            desktopExitRunnable?.invoke(transition)
         }
-        exitResult.asExit()?.runOnTransitionStart?.invoke(transition)
-        desktopExitRunnable?.invoke(transition)
     }
 
     /** Move a task with given `taskId` to fullscreen */
@@ -1845,7 +1852,11 @@ class DesktopTasksController(
         displayId: Int,
         forceExitDesktop: Boolean,
     ): Boolean {
-        if (forceExitDesktop && DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
+        if (
+            forceExitDesktop &&
+                (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue ||
+                    DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PIP.isTrue)
+        ) {
             // |forceExitDesktop| is true when the callers knows we'll exit desktop, such as when
             // explicitly going fullscreen, so there's no point in checking the desktop state.
             return true
@@ -1862,6 +1873,33 @@ class DesktopTasksController(
         return true
     }
 
+    /** Potentially perform Desktop cleanup after a task successfully enters PiP. */
+    @VisibleForTesting
+    fun onDesktopTaskEnteredPip(
+        taskId: Int,
+        deskId: Int,
+        displayId: Int,
+        taskIsLastVisibleTaskBeforePip: Boolean,
+    ) {
+        if (
+            !willExitDesktop(taskId, displayId, forceExitDesktop = taskIsLastVisibleTaskBeforePip)
+        ) {
+            return
+        }
+
+        val wct = WindowContainerTransaction()
+        val desktopExitRunnable =
+            performDesktopExitCleanUp(
+                wct = wct,
+                deskId = deskId,
+                displayId = displayId,
+                willExitDesktop = true,
+            )
+
+        val transition = transitions.startTransition(TRANSIT_CHANGE, wct, /* handler= */ null)
+        desktopExitRunnable?.invoke(transition)
+    }
+
     private fun performDesktopExitCleanupIfNeeded(
         taskId: Int,
         deskId: Int? = null,
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 2ec6105e5af9..df4d18f8c803 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
@@ -23,7 +23,6 @@ import android.os.IBinder
 import android.view.SurfaceControl
 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.window.DesktopExperienceFlags
 import android.window.DesktopModeFlags
@@ -43,8 +42,7 @@ import com.android.wm.shell.shared.TransitionUtil.isOpeningMode
 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 java.util.Optional
 
 /**
  * A [Transitions.TransitionObserver] that observes shell transitions and updates the
@@ -57,6 +55,7 @@ class DesktopTasksTransitionObserver(
     private val transitions: Transitions,
     private val shellTaskOrganizer: ShellTaskOrganizer,
     private val desktopMixedTransitionHandler: DesktopMixedTransitionHandler,
+    private val desktopPipTransitionObserver: Optional<DesktopPipTransitionObserver>,
     private val backAnimationController: BackAnimationController,
     private val desktopWallpaperActivityTokenProvider: DesktopWallpaperActivityTokenProvider,
     shellInit: ShellInit,
@@ -65,8 +64,6 @@ class DesktopTasksTransitionObserver(
     data class CloseWallpaperTransition(val transition: IBinder, val displayId: Int)
 
     private var transitionToCloseWallpaper: CloseWallpaperTransition? = null
-    /* Pending PiP transition and its associated display id and task id. */
-    private var pendingPipTransitionAndPipTask: Triple<IBinder, Int, Int>? = null
     private var currentProfileId: Int
 
     init {
@@ -100,33 +97,7 @@ class DesktopTasksTransitionObserver(
             removeTaskIfNeeded(info)
         }
         removeWallpaperOnLastTaskClosingIfNeeded(transition, info)
-
-        val desktopRepository = desktopUserRepositories.getProfile(currentProfileId)
-        info.changes.forEach { change ->
-            change.taskInfo?.let { taskInfo ->
-                if (
-                    DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PIP.isTrue &&
-                        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,
-                            )
-                    }
-                }
-            }
-        }
+        desktopPipTransitionObserver.ifPresent { it.onTransitionReady(transition, info) }
     }
 
     private fun removeTaskIfNeeded(info: TransitionInfo) {
@@ -301,18 +272,6 @@ 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
         }
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopPipTransitionObserverTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopPipTransitionObserverTest.kt
new file mode 100644
index 000000000000..ef394d81cc57
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopPipTransitionObserverTest.kt
@@ -0,0 +1,98 @@
+/*
+ * Copyright (C) 2025 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
+
+import android.app.WindowConfiguration.WINDOWING_MODE_PINNED
+import android.os.Binder
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import android.testing.AndroidTestingRunner
+import android.view.WindowManager.TRANSIT_PIP
+import android.window.TransitionInfo
+import androidx.test.filters.SmallTest
+import com.android.window.flags.Flags
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.TestRunningTaskInfoBuilder
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+
+/**
+ * Tests for [DesktopPipTransitionObserver].
+ *
+ * Build/Install/Run: atest WMShellUnitTests:DesktopPipTransitionObserverTest
+ */
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class DesktopPipTransitionObserverTest : ShellTestCase() {
+
+    @JvmField @Rule val setFlagsRule = SetFlagsRule()
+
+    private lateinit var observer: DesktopPipTransitionObserver
+
+    private val transition = Binder()
+    private var onSuccessInvokedCount = 0
+
+    @Before
+    fun setUp() {
+        observer = DesktopPipTransitionObserver()
+
+        onSuccessInvokedCount = 0
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun onTransitionReady_taskInPinnedWindowingMode_onSuccessInvoked() {
+        val taskId = 1
+        val pipTransition = createPendingPipTransition(taskId)
+        val successfulChange = createChange(taskId, WINDOWING_MODE_PINNED)
+        observer.addPendingPipTransition(pipTransition)
+
+        observer.onTransitionReady(
+            transition = transition,
+            info = TransitionInfo(
+                TRANSIT_PIP, /* flags= */
+                0
+            ).apply { addChange(successfulChange) },
+        )
+
+        assertThat(onSuccessInvokedCount).isEqualTo(1)
+    }
+
+    private fun createPendingPipTransition(
+        taskId: Int
+    ): DesktopPipTransitionObserver.PendingPipTransition {
+        return DesktopPipTransitionObserver.PendingPipTransition(
+            token = transition,
+            taskId = taskId,
+            onSuccess = { onSuccessInvokedCount += 1 },
+        )
+    }
+
+    private fun createChange(taskId: Int, windowingMode: Int): TransitionInfo.Change {
+        return TransitionInfo.Change(mock(), mock()).apply {
+            taskInfo =
+                TestRunningTaskInfoBuilder()
+                    .setTaskId(taskId)
+                    .setWindowingMode(windowingMode)
+                    .build()
+        }
+    }
+}
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 c455205f6411..a10aeca95bce 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
@@ -1237,36 +1237,6 @@ class DesktopRepositoryTest(flags: FlagsParameterization) : ShellTestCase() {
     }
 
     @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.addDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
-        repo.setActiveDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
-        repo.setTaskInPip(DEFAULT_DESKTOP_ID, taskId = 1, enterPip = true)
-        repo.setTaskInPip(SECOND_DISPLAY, taskId = 2, enterPip = true)
-
-        assertThat(repo.isTaskMinimizedPipInDisplay(DEFAULT_DESKTOP_ID, taskId = 1)).isTrue()
-        assertThat(repo.isTaskMinimizedPipInDisplay(SECOND_DISPLAY, taskId = 2)).isTrue()
-    }
-
-    @Test
     @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun addTask_deskDoesNotExists_createsDesk() {
         repo.addTask(displayId = 999, taskId = 6, isVisible = true)
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 8d8140c47c1a..7efcd4fc3c8f 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
@@ -265,6 +265,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     @Mock private lateinit var desksOrganizer: DesksOrganizer
     @Mock private lateinit var userProfileContexts: UserProfileContexts
     @Mock private lateinit var desksTransitionsObserver: DesksTransitionObserver
+    @Mock private lateinit var desktopPipTransitionObserver: DesktopPipTransitionObserver
     @Mock private lateinit var packageManager: PackageManager
     @Mock private lateinit var mockDisplayContext: Context
     @Mock private lateinit var dragToDisplayTransitionHandler: DragToDisplayTransitionHandler
@@ -393,6 +394,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(wallpaperToken)
         whenever(userProfileContexts[anyInt()]).thenReturn(context)
         whenever(userProfileContexts.getOrCreate(anyInt())).thenReturn(context)
+        whenever(freeformTaskTransitionStarter.startPipTransition(any())).thenReturn(Binder())
 
         controller = createController()
         controller.setSplitScreenController(splitScreenController)
@@ -457,6 +459,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
             overviewToDesktopTransitionObserver,
             desksOrganizer,
             desksTransitionsObserver,
+            Optional.of(desktopPipTransitionObserver),
             userProfileContexts,
             desktopModeCompatPolicy,
             dragToDisplayTransitionHandler,
@@ -3502,6 +3505,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
     fun onPipTaskMinimize_autoEnterEnabled_startPipTransition() {
         val task = setUpPipTask(autoEnterEnabled = true)
         val handler = mock(TransitionHandler::class.java)
@@ -3516,6 +3520,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
     fun onPipTaskMinimize_autoEnterDisabled_startMinimizeTransition() {
         val task = setUpPipTask(autoEnterEnabled = false)
         whenever(
@@ -3535,6 +3540,90 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WALLPAPER_ACTIVITY_FOR_SYSTEM_USER,
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP,
+    )
+    fun onDesktopTaskEnteredPip_pipIsLastTask_removesWallpaper() {
+        val task = setUpPipTask(autoEnterEnabled = true)
+
+        controller.onDesktopTaskEnteredPip(
+            taskId = task.taskId,
+            deskId = DEFAULT_DISPLAY,
+            displayId = task.displayId,
+            taskIsLastVisibleTaskBeforePip = true,
+        )
+
+        // Wallpaper is moved to the back
+        val wct = getLatestTransition()
+        wct.assertReorder(wallpaperToken, /* toTop= */ false)
+    }
+
+    @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun onDesktopTaskEnteredPip_pipIsLastTask_deactivatesDesk() {
+        val deskId = DEFAULT_DISPLAY
+        val task = setUpPipTask(autoEnterEnabled = true, deskId = deskId)
+        val transition = Binder()
+        whenever(transitions.startTransition(any(), any(), anyOrNull())).thenReturn(transition)
+
+        controller.onDesktopTaskEnteredPip(
+            taskId = task.taskId,
+            deskId = deskId,
+            displayId = task.displayId,
+            taskIsLastVisibleTaskBeforePip = true,
+        )
+
+        verify(desksOrganizer).deactivateDesk(any(), eq(deskId))
+        verify(desksTransitionsObserver)
+            .addPendingTransition(DeskTransition.DeactivateDesk(transition, deskId))
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun onDesktopTaskEnteredPip_pipIsLastTask_launchesHome() {
+        val task = setUpPipTask(autoEnterEnabled = true)
+
+        controller.onDesktopTaskEnteredPip(
+            taskId = task.taskId,
+            deskId = DEFAULT_DISPLAY,
+            displayId = task.displayId,
+            taskIsLastVisibleTaskBeforePip = true,
+        )
+
+        val wct = getLatestTransition()
+        wct.assertPendingIntent(launchHomeIntent(DEFAULT_DISPLAY))
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+    fun onDesktopTaskEnteredPip_pipIsNotLastTask_doesntExitDesktopMode() {
+        val task = setUpPipTask(autoEnterEnabled = true)
+        val deskId = DEFAULT_DISPLAY
+        setUpFreeformTask(deskId = deskId) // launch another freeform task
+        val transition = Binder()
+        whenever(transitions.startTransition(any(), any(), anyOrNull())).thenReturn(transition)
+
+        controller.onDesktopTaskEnteredPip(
+            taskId = task.taskId,
+            deskId = deskId,
+            displayId = task.displayId,
+            taskIsLastVisibleTaskBeforePip = false,
+        )
+
+        // No transition to exit Desktop mode is started
+        verifyWCTNotExecuted()
+        verify(desktopModeEnterExitTransitionListener, never())
+            .onExitDesktopModeTransitionStarted(FULLSCREEN_ANIMATION_DURATION)
+        verify(desksOrganizer, never()).deactivateDesk(any(), eq(deskId))
+        verify(desksTransitionsObserver, never())
+            .addPendingTransition(DeskTransition.DeactivateDesk(transition, deskId))
+    }
+
+    @Test
     fun onDesktopWindowMinimize_singleActiveTask_noWallpaperActivityToken_doesntRemoveWallpaper() {
         val task = setUpFreeformTask(active = true)
         val transition = Binder()
@@ -7656,8 +7745,12 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         return task
     }
 
-    private fun setUpPipTask(autoEnterEnabled: Boolean): RunningTaskInfo =
-        setUpFreeformTask().apply {
+    private fun setUpPipTask(
+        autoEnterEnabled: Boolean,
+        displayId: Int = DEFAULT_DISPLAY,
+        deskId: Int = DEFAULT_DISPLAY,
+    ): RunningTaskInfo =
+        setUpFreeformTask(displayId = displayId, deskId = deskId).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 ac9a509ac6cb..5ef1ace7873d 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,7 +22,6 @@ 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.DisableFlags
 import android.platform.test.annotations.EnableFlags
@@ -30,7 +29,6 @@ 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
@@ -41,7 +39,6 @@ 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
@@ -51,10 +48,9 @@ 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 java.util.Optional
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -90,6 +86,7 @@ class DesktopTasksTransitionObserverTest {
     private val userRepositories = mock<DesktopUserRepositories>()
     private val taskRepository = mock<DesktopRepository>()
     private val mixedHandler = mock<DesktopMixedTransitionHandler>()
+    private val pipTransitionObserver = mock<DesktopPipTransitionObserver>()
     private val backAnimationController = mock<BackAnimationController>()
     private val desktopWallpaperActivityTokenProvider =
         mock<DesktopWallpaperActivityTokenProvider>()
@@ -114,6 +111,7 @@ class DesktopTasksTransitionObserverTest {
                 transitions,
                 shellTaskOrganizer,
                 mixedHandler,
+                Optional.of(pipTransitionObserver),
                 backAnimationController,
                 desktopWallpaperActivityTokenProvider,
                 shellInit,
@@ -393,56 +391,6 @@ class DesktopTasksTransitionObserverTest {
         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, type = 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,