[PiP on Desktop] Add PipDesktopState for PiP checks in Desktop Mode.

This is a refactoring work to provide a centralized place to check PiP
in Desktop Mode status.

Bug: 350475854
Test: atest DesktopTasksTransitionObserverTest DesktopRepositoryTest
DesktopTasksControllerTest PipSchedulerTest PipDesktopStateTest
Test: Manually verify all PiP on Desktop CUJs are WAI.
Flag: com.android.window.flags.enable_desktop_windowing_pip

Change-Id: I06114e0bb4194fefa8a38cc2b4fc6c94c47f3907

6 files changed, 403 insertions, 145 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipDesktopState.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipDesktopState.java
new file mode 100644
index 000000000000..c10c2c905c97
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipDesktopState.java
@@ -0,0 +1,159 @@
+/*
+ * 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.common.pip;
+
+import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
+import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
+
+import android.app.ActivityManager;
+import android.window.DisplayAreaInfo;
+import android.window.WindowContainerToken;
+import android.window.WindowContainerTransaction;
+
+import com.android.window.flags.Flags;
+import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
+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.pip2.phone.PipTransition;
+
+import java.util.Optional;
+
+/** Helper class for PiP on Desktop Mode. */
+public class PipDesktopState {
+    private final PipDisplayLayoutState mPipDisplayLayoutState;
+    private final Optional<DesktopUserRepositories> mDesktopUserRepositoriesOptional;
+    private final Optional<DesktopWallpaperActivityTokenProvider>
+            mDesktopWallpaperActivityTokenProviderOptional;
+    private final RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
+
+    public PipDesktopState(PipDisplayLayoutState pipDisplayLayoutState,
+            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
+            Optional<DesktopWallpaperActivityTokenProvider>
+                    desktopWallpaperActivityTokenProviderOptional,
+            RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
+        mPipDisplayLayoutState = pipDisplayLayoutState;
+        mDesktopUserRepositoriesOptional = desktopUserRepositoriesOptional;
+        mDesktopWallpaperActivityTokenProviderOptional =
+                desktopWallpaperActivityTokenProviderOptional;
+        mRootTaskDisplayAreaOrganizer = rootTaskDisplayAreaOrganizer;
+    }
+
+    /**
+     * Returns whether PiP in Desktop Windowing is enabled by checking the following:
+     * - Desktop Windowing in PiP flag is enabled
+     * - DesktopWallpaperActivityTokenProvider is injected
+     * - DesktopUserRepositories is injected
+     */
+    public boolean isDesktopWindowingPipEnabled() {
+        return Flags.enableDesktopWindowingPip()
+                && mDesktopWallpaperActivityTokenProviderOptional.isPresent()
+                && mDesktopUserRepositoriesOptional.isPresent();
+    }
+
+    /** Returns whether PiP in Connected Displays is enabled by checking the flag. */
+    public boolean isConnectedDisplaysPipEnabled() {
+        return Flags.enableConnectedDisplaysPip();
+    }
+
+    /** Returns whether the display with the PiP task is in freeform windowing mode. */
+    private boolean isDisplayInFreeform() {
+        final DisplayAreaInfo tdaInfo = mRootTaskDisplayAreaOrganizer.getDisplayAreaInfo(
+                mPipDisplayLayoutState.getDisplayId());
+        if (tdaInfo != null) {
+            return tdaInfo.configuration.windowConfiguration.getWindowingMode()
+                    == WINDOWING_MODE_FREEFORM;
+        }
+        return false;
+    }
+
+    /** Returns whether PiP is exiting while we're in a Desktop Mode session. */
+    private boolean isPipExitingToDesktopMode() {
+        // Early return if PiP in Desktop Windowing is not supported.
+        if (!isDesktopWindowingPipEnabled()) {
+            return false;
+        }
+        final int displayId = mPipDisplayLayoutState.getDisplayId();
+        return getDesktopRepository().getVisibleTaskCount(displayId) > 0
+                || getDesktopWallpaperActivityTokenProvider().isWallpaperActivityVisible(displayId)
+                || isDisplayInFreeform();
+    }
+
+    /** Returns whether {@param pipTask} would be entering in a Desktop Mode session. */
+    public boolean isPipEnteringInDesktopMode(ActivityManager.RunningTaskInfo pipTask) {
+        // Early return if PiP in Desktop Windowing is not supported.
+        if (!isDesktopWindowingPipEnabled()) {
+            return false;
+        }
+        final DesktopRepository desktopRepository = getDesktopRepository();
+        return desktopRepository.getVisibleTaskCount(pipTask.getDisplayId()) > 0
+                || desktopRepository.isMinimizedPipPresentInDisplay(pipTask.getDisplayId());
+    }
+
+    /**
+     * Invoked when an EXIT_PiP transition is detected in {@link PipTransition}.
+     * Returns whether the PiP exiting should also trigger the active Desktop Mode session to exit.
+     */
+    public boolean shouldExitPipExitDesktopMode() {
+        // Early return if PiP in Desktop Windowing is not supported.
+        if (!isDesktopWindowingPipEnabled()) {
+            return false;
+        }
+        final int displayId = mPipDisplayLayoutState.getDisplayId();
+        return getDesktopRepository().getVisibleTaskCount(displayId) == 0
+                && getDesktopWallpaperActivityTokenProvider().isWallpaperActivityVisible(displayId);
+    }
+
+    /**
+     * Returns a {@link WindowContainerTransaction} that reorders the {@link WindowContainerToken}
+     * of the DesktopWallpaperActivity for the display with the given {@param displayId}.
+     */
+    public WindowContainerTransaction getWallpaperActivityTokenWct(int displayId) {
+        return new WindowContainerTransaction().reorder(
+                getDesktopWallpaperActivityTokenProvider().getToken(displayId), /* onTop= */ false);
+    }
+
+    /**
+     * The windowing mode to restore to when resizing out of PIP direction.
+     * Defaults to undefined and can be overridden to restore to an alternate windowing mode.
+     */
+    public int getOutPipWindowingMode() {
+        // If we are exiting PiP while the device is in Desktop mode (the task should expand to
+        // freeform windowing mode):
+        // 1) If the display windowing mode is freeform, set windowing mode to UNDEFINED so it will
+        //    resolve the windowing mode to the display's windowing mode.
+        // 2) If the display windowing mode is not FREEFORM, set windowing mode to FREEFORM.
+        if (isPipExitingToDesktopMode()) {
+            if (isDisplayInFreeform()) {
+                return WINDOWING_MODE_UNDEFINED;
+            } else {
+                return WINDOWING_MODE_FREEFORM;
+            }
+        }
+
+        // By default, or if the task is going to fullscreen, reset the windowing mode to undefined.
+        return WINDOWING_MODE_UNDEFINED;
+    }
+
+    private DesktopRepository getDesktopRepository() {
+        return mDesktopUserRepositoriesOptional.get().getCurrent();
+    }
+
+    private DesktopWallpaperActivityTokenProvider getDesktopWallpaperActivityTokenProvider() {
+        return mDesktopWallpaperActivityTokenProviderOptional.get();
+    }
+}
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 413300612f7d..e7c76bbd91b2 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
@@ -31,6 +31,7 @@ import com.android.wm.shell.common.TaskStackListenerImpl;
 import com.android.wm.shell.common.pip.PipAppOpsListener;
 import com.android.wm.shell.common.pip.PipBoundsAlgorithm;
 import com.android.wm.shell.common.pip.PipBoundsState;
+import com.android.wm.shell.common.pip.PipDesktopState;
 import com.android.wm.shell.common.pip.PipDisplayLayoutState;
 import com.android.wm.shell.common.pip.PipMediaController;
 import com.android.wm.shell.common.pip.PipPerfHintController;
@@ -84,14 +85,11 @@ public abstract class Pip2Module {
             @NonNull PipDisplayLayoutState pipDisplayLayoutState,
             @NonNull PipUiStateChangeController pipUiStateChangeController,
             DisplayController displayController,
-            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
-            Optional<DesktopWallpaperActivityTokenProvider>
-                    desktopWallpaperActivityTokenProviderOptional) {
+            PipDesktopState pipDesktopState) {
         return new PipTransition(context, shellInit, shellTaskOrganizer, transitions,
                 pipBoundsState, null, pipBoundsAlgorithm, pipTaskListener,
                 pipScheduler, pipStackListenerController, pipDisplayLayoutState,
-                pipUiStateChangeController, displayController, desktopUserRepositoriesOptional,
-                desktopWallpaperActivityTokenProviderOptional);
+                pipUiStateChangeController, displayController, pipDesktopState);
     }
 
     @WMSingleton
@@ -142,13 +140,9 @@ public abstract class Pip2Module {
             PipBoundsState pipBoundsState,
             @ShellMainThread ShellExecutor mainExecutor,
             PipTransitionState pipTransitionState,
-            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
-            Optional<DesktopWallpaperActivityTokenProvider>
-                    desktopWallpaperActivityTokenProviderOptional,
-            RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
+            PipDesktopState pipDesktopState) {
         return new PipScheduler(context, pipBoundsState, mainExecutor, pipTransitionState,
-                desktopUserRepositoriesOptional, desktopWallpaperActivityTokenProviderOptional,
-                rootTaskDisplayAreaOrganizer);
+                pipDesktopState);
     }
 
     @WMSingleton
@@ -233,4 +227,17 @@ public abstract class Pip2Module {
         return new PipTaskListener(context, shellTaskOrganizer, pipTransitionState,
                 pipScheduler, pipBoundsState, pipBoundsAlgorithm, mainExecutor);
     }
+
+    @WMSingleton
+    @Provides
+    static PipDesktopState providePipDesktopState(
+            PipDisplayLayoutState pipDisplayLayoutState,
+            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
+            Optional<DesktopWallpaperActivityTokenProvider>
+                    desktopWallpaperActivityTokenProviderOptional,
+            RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer
+    ) {
+        return new PipDesktopState(pipDisplayLayoutState, desktopUserRepositoriesOptional,
+                desktopWallpaperActivityTokenProviderOptional, rootTaskDisplayAreaOrganizer);
+    }
 }
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 21b0820f523a..e17587ff18bc 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
@@ -16,14 +16,10 @@
 
 package com.android.wm.shell.pip2.phone;
 
-import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
-import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
-
 import android.content.Context;
 import android.graphics.Matrix;
 import android.graphics.Rect;
 import android.view.SurfaceControl;
-import android.window.DisplayAreaInfo;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
@@ -32,20 +28,14 @@ import androidx.annotation.Nullable;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.protolog.ProtoLog;
-import com.android.window.flags.Flags;
-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.common.pip.PipDesktopState;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.pip2.PipSurfaceTransactionHelper;
 import com.android.wm.shell.pip2.animation.PipAlphaAnimator;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 
-import java.util.Objects;
-import java.util.Optional;
-
 /**
  * Scheduler for Shell initiated PiP transitions and animations.
  */
@@ -56,10 +46,7 @@ public class PipScheduler {
     private final PipBoundsState mPipBoundsState;
     private final ShellExecutor mMainExecutor;
     private final PipTransitionState mPipTransitionState;
-    private final Optional<DesktopUserRepositories> mDesktopUserRepositoriesOptional;
-    private final Optional<DesktopWallpaperActivityTokenProvider>
-            mDesktopWallpaperActivityTokenProviderOptional;
-    private final RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
+    private final PipDesktopState mPipDesktopState;
     private PipTransitionController mPipTransitionController;
     private PipSurfaceTransactionHelper.SurfaceControlTransactionFactory
             mSurfaceControlTransactionFactory;
@@ -72,18 +59,12 @@ public class PipScheduler {
             PipBoundsState pipBoundsState,
             ShellExecutor mainExecutor,
             PipTransitionState pipTransitionState,
-            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
-            Optional<DesktopWallpaperActivityTokenProvider>
-                    desktopWallpaperActivityTokenProviderOptional,
-            RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer) {
+            PipDesktopState pipDesktopState) {
         mContext = context;
         mPipBoundsState = pipBoundsState;
         mMainExecutor = mainExecutor;
         mPipTransitionState = pipTransitionState;
-        mDesktopUserRepositoriesOptional = desktopUserRepositoriesOptional;
-        mDesktopWallpaperActivityTokenProviderOptional =
-                desktopWallpaperActivityTokenProviderOptional;
-        mRootTaskDisplayAreaOrganizer = rootTaskDisplayAreaOrganizer;
+        mPipDesktopState = pipDesktopState;
 
         mSurfaceControlTransactionFactory =
                 new PipSurfaceTransactionHelper.VsyncSurfaceControlTransactionFactory();
@@ -105,7 +86,7 @@ public class PipScheduler {
         wct.setBounds(pipTaskToken, null);
         // if we are hitting a multi-activity case
         // windowing mode change will reparent to original host task
-        wct.setWindowingMode(pipTaskToken, getOutPipWindowingMode());
+        wct.setWindowingMode(pipTaskToken, mPipDesktopState.getOutPipWindowingMode());
         return wct;
     }
 
@@ -235,55 +216,6 @@ public class PipScheduler {
         maybeUpdateMovementBounds();
     }
 
-    /** Returns whether the display is in freeform windowing mode. */
-    private boolean isDisplayInFreeform() {
-        final DisplayAreaInfo tdaInfo = mRootTaskDisplayAreaOrganizer.getDisplayAreaInfo(
-                Objects.requireNonNull(mPipTransitionState.getPipTaskInfo()).displayId);
-        if (tdaInfo != null) {
-            return tdaInfo.configuration.windowConfiguration.getWindowingMode()
-                    == WINDOWING_MODE_FREEFORM;
-        }
-        return false;
-    }
-
-    /** Returns whether PiP is exiting while we're in desktop mode. */
-    private boolean isPipExitingToDesktopMode() {
-        // 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();
-    }
-
-    /**
-     * The windowing mode to restore to when resizing out of PIP direction. Defaults to undefined
-     * and can be overridden to restore to an alternate windowing mode.
-     */
-    private int getOutPipWindowingMode() {
-        // If we are exiting PiP while the device is in Desktop mode (the task should expand to
-        // freeform windowing mode):
-        // 1) If the display windowing mode is freeform, set windowing mode to undefined so it will
-        //    resolve the windowing mode to the display's windowing mode.
-        // 2) If the display windowing mode is not freeform, set windowing mode to freeform.
-        if (isPipExitingToDesktopMode()) {
-            if (isDisplayInFreeform()) {
-                return WINDOWING_MODE_UNDEFINED;
-            } else {
-                return WINDOWING_MODE_FREEFORM;
-            }
-        }
-
-        // By default, or if the task is going to fullscreen, reset the windowing mode to undefined.
-        return WINDOWING_MODE_UNDEFINED;
-    }
-
     @VisibleForTesting
     void setSurfaceControlTransactionFactory(
             @NonNull PipSurfaceTransactionHelper.SurfaceControlTransactionFactory factory) {
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 229962488acf..d3e630ddc703 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
@@ -56,19 +56,16 @@ import android.window.WindowContainerTransaction;
 import androidx.annotation.Nullable;
 
 import com.android.internal.util.Preconditions;
-import com.android.window.flags.Flags;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.ComponentUtils;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.pip.PipBoundsAlgorithm;
 import com.android.wm.shell.common.pip.PipBoundsState;
+import com.android.wm.shell.common.pip.PipDesktopState;
 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.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;
@@ -79,8 +76,6 @@ import com.android.wm.shell.shared.pip.PipContentOverlay;
 import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
 
-import java.util.Optional;
-
 /**
  * Implementation of transitions for PiP on phone.
  */
@@ -117,9 +112,7 @@ public class PipTransition extends PipTransitionController implements
     private final PipDisplayLayoutState mPipDisplayLayoutState;
     private final DisplayController mDisplayController;
     private final PipSurfaceTransactionHelper mPipSurfaceTransactionHelper;
-    private final Optional<DesktopUserRepositories> mDesktopUserRepositoriesOptional;
-    private final Optional<DesktopWallpaperActivityTokenProvider>
-            mDesktopWallpaperActivityTokenProviderOptional;
+    private final PipDesktopState mPipDesktopState;
 
     //
     // Transition caches
@@ -159,9 +152,7 @@ public class PipTransition extends PipTransitionController implements
             PipDisplayLayoutState pipDisplayLayoutState,
             PipUiStateChangeController pipUiStateChangeController,
             DisplayController displayController,
-            Optional<DesktopUserRepositories> desktopUserRepositoriesOptional,
-            Optional<DesktopWallpaperActivityTokenProvider>
-                    desktopWallpaperActivityTokenProviderOptional) {
+            PipDesktopState pipDesktopState) {
         super(shellInit, shellTaskOrganizer, transitions, pipBoundsState, pipMenuController,
                 pipBoundsAlgorithm);
 
@@ -174,9 +165,7 @@ public class PipTransition extends PipTransitionController implements
         mPipDisplayLayoutState = pipDisplayLayoutState;
         mDisplayController = displayController;
         mPipSurfaceTransactionHelper = new PipSurfaceTransactionHelper(mContext);
-        mDesktopUserRepositoriesOptional = desktopUserRepositoriesOptional;
-        mDesktopWallpaperActivityTokenProviderOptional =
-                desktopWallpaperActivityTokenProviderOptional;
+        mPipDesktopState = pipDesktopState;
     }
 
     @Override
@@ -800,7 +789,7 @@ public class PipTransition extends PipTransitionController implements
                     && getFixedRotationDelta(info, pipTaskChange) == ROTATION_90) {
                 adjustedSourceRectHint.offset(cutoutInsets.left, cutoutInsets.top);
             }
-            if (Flags.enableDesktopWindowingPip()) {
+            if (mPipDesktopState.isDesktopWindowingPipEnabled()) {
                 adjustedSourceRectHint.offset(-pipActivityChange.getStartAbsBounds().left,
                         -pipActivityChange.getStartAbsBounds().top);
             }
@@ -860,7 +849,7 @@ public class PipTransition extends PipTransitionController implements
 
         // If PiP is enabled on Connected Displays, update PipDisplayLayoutState to have the correct
         // display info that PiP is entering in.
-        if (Flags.enableConnectedDisplaysPip()) {
+        if (mPipDesktopState.isConnectedDisplaysPipEnabled()) {
             final DisplayLayout displayLayout = mDisplayController.getDisplayLayout(
                     pipTask.displayId);
             if (displayLayout != null) {
@@ -908,12 +897,7 @@ public class PipTransition extends PipTransitionController implements
         // 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()
-                        .isMinimizedPipPresentInDisplay(pipTask.displayId));
-        if (isInDesktopSession) {
+        if (mPipDesktopState.isPipEnteringInDesktopMode(pipTask)) {
             return false;
         }
 
@@ -1087,26 +1071,13 @@ 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
-                        );
-                    }
+                if (mPipDesktopState.shouldExitPipExitDesktopMode()) {
+                    mTransitions.startTransition(
+                            TRANSIT_TO_BACK,
+                            mPipDesktopState.getWallpaperActivityTokenWct(
+                                    mPipTransitionState.getPipTaskInfo().getDisplayId()),
+                            null /* firstHandler */
+                    );
                 }
                 mPipTransitionState.setPinnedTaskLeash(null);
                 mPipTransitionState.setPipTaskInfo(null);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/pip/PipDesktopStateTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/pip/PipDesktopStateTest.java
new file mode 100644
index 000000000000..75ad621e1cad
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/pip/PipDesktopStateTest.java
@@ -0,0 +1,204 @@
+/*
+ * 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.common.pip;
+
+import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
+import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
+import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
+
+import static com.android.window.flags.Flags.FLAG_ENABLE_CONNECTED_DISPLAYS_PIP;
+import static com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PIP;
+
+import static junit.framework.Assert.assertEquals;
+import static junit.framework.Assert.assertFalse;
+import static junit.framework.Assert.assertTrue;
+
+import static org.mockito.Mockito.when;
+
+import android.app.ActivityManager;
+import android.platform.test.annotations.EnableFlags;
+import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
+import android.window.DisplayAreaInfo;
+import android.window.WindowContainerToken;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
+import com.android.wm.shell.desktopmode.DesktopRepository;
+import com.android.wm.shell.desktopmode.DesktopUserRepositories;
+import com.android.wm.shell.desktopmode.desktopwallpaperactivity.DesktopWallpaperActivityTokenProvider;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+
+import java.util.Optional;
+
+/**
+ * Unit test against {@link PipDesktopState}.
+ */
+@SmallTest
+@TestableLooper.RunWithLooper
+@RunWith(AndroidTestingRunner.class)
+@EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PIP)
+public class PipDesktopStateTest {
+    @Mock private PipDisplayLayoutState mMockPipDisplayLayoutState;
+    @Mock private Optional<DesktopUserRepositories> mMockDesktopUserRepositoriesOptional;
+    @Mock private Optional<DesktopWallpaperActivityTokenProvider>
+            mMockDesktopWallpaperActivityTokenProviderOptional;
+    @Mock private DesktopUserRepositories mMockDesktopUserRepositories;
+    @Mock private DesktopWallpaperActivityTokenProvider mMockDesktopWallpaperActivityTokenProvider;
+    @Mock private DesktopRepository mMockDesktopRepository;
+    @Mock private RootTaskDisplayAreaOrganizer mMockRootTaskDisplayAreaOrganizer;
+    @Mock private ActivityManager.RunningTaskInfo mMockTaskInfo;
+
+    private static final int DISPLAY_ID = 1;
+    private DisplayAreaInfo mDefaultTda;
+    private PipDesktopState mPipDesktopState;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        when(mMockDesktopUserRepositoriesOptional.get()).thenReturn(mMockDesktopUserRepositories);
+        when(mMockDesktopWallpaperActivityTokenProviderOptional.get()).thenReturn(
+                mMockDesktopWallpaperActivityTokenProvider);
+        when(mMockDesktopUserRepositories.getCurrent()).thenReturn(mMockDesktopRepository);
+        when(mMockDesktopUserRepositoriesOptional.isPresent()).thenReturn(true);
+        when(mMockDesktopWallpaperActivityTokenProviderOptional.isPresent()).thenReturn(true);
+
+        when(mMockTaskInfo.getDisplayId()).thenReturn(DISPLAY_ID);
+        when(mMockPipDisplayLayoutState.getDisplayId()).thenReturn(DISPLAY_ID);
+
+        mDefaultTda = new DisplayAreaInfo(Mockito.mock(WindowContainerToken.class), DISPLAY_ID, 0);
+        when(mMockRootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DISPLAY_ID)).thenReturn(
+                mDefaultTda);
+
+        mPipDesktopState = new PipDesktopState(mMockPipDisplayLayoutState,
+                mMockDesktopUserRepositoriesOptional,
+                mMockDesktopWallpaperActivityTokenProviderOptional,
+                mMockRootTaskDisplayAreaOrganizer);
+    }
+
+    @Test
+    public void isDesktopWindowingPipEnabled_returnsTrue() {
+        assertTrue(mPipDesktopState.isDesktopWindowingPipEnabled());
+    }
+
+    @Test
+    public void isDesktopWindowingPipEnabled_desktopRepositoryEmpty_returnsFalse() {
+        when(mMockDesktopUserRepositoriesOptional.isPresent()).thenReturn(false);
+
+        assertFalse(mPipDesktopState.isDesktopWindowingPipEnabled());
+    }
+
+    @Test
+    public void isDesktopWindowingPipEnabled_desktopWallpaperEmpty_returnsFalse() {
+        when(mMockDesktopWallpaperActivityTokenProviderOptional.isPresent()).thenReturn(false);
+
+        assertFalse(mPipDesktopState.isDesktopWindowingPipEnabled());
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_CONNECTED_DISPLAYS_PIP)
+    public void isConnectedDisplaysPipEnabled_returnsTrue() {
+        assertTrue(mPipDesktopState.isConnectedDisplaysPipEnabled());
+    }
+
+    @Test
+    public void isPipEnteringInDesktopMode_visibleCountZero_minimizedPipPresent_returnsTrue() {
+        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mMockDesktopRepository.isMinimizedPipPresentInDisplay(DISPLAY_ID)).thenReturn(true);
+
+        assertTrue(mPipDesktopState.isPipEnteringInDesktopMode(mMockTaskInfo));
+    }
+
+    @Test
+    public void isPipEnteringInDesktopMode_visibleCountNonzero_minimizedPipAbsent_returnsTrue() {
+        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(1);
+        when(mMockDesktopRepository.isMinimizedPipPresentInDisplay(DISPLAY_ID)).thenReturn(false);
+
+        assertTrue(mPipDesktopState.isPipEnteringInDesktopMode(mMockTaskInfo));
+    }
+
+    @Test
+    public void isPipEnteringInDesktopMode_visibleCountZero_minimizedPipAbsent_returnsFalse() {
+        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mMockDesktopRepository.isMinimizedPipPresentInDisplay(DISPLAY_ID)).thenReturn(false);
+
+        assertFalse(mPipDesktopState.isPipEnteringInDesktopMode(mMockTaskInfo));
+    }
+
+    @Test
+    public void shouldExitPipExitDesktopMode_visibleCountZero_wallpaperInvisible_returnsFalse() {
+        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mMockDesktopWallpaperActivityTokenProvider.isWallpaperActivityVisible(
+                DISPLAY_ID)).thenReturn(false);
+
+        assertFalse(mPipDesktopState.shouldExitPipExitDesktopMode());
+    }
+
+    @Test
+    public void shouldExitPipExitDesktopMode_visibleCountNonzero_wallpaperVisible_returnsFalse() {
+        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(1);
+        when(mMockDesktopWallpaperActivityTokenProvider.isWallpaperActivityVisible(
+                DISPLAY_ID)).thenReturn(true);
+
+        assertFalse(mPipDesktopState.shouldExitPipExitDesktopMode());
+    }
+
+    @Test
+    public void shouldExitPipExitDesktopMode_visibleCountZero_wallpaperVisible_returnsTrue() {
+        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mMockDesktopWallpaperActivityTokenProvider.isWallpaperActivityVisible(
+                DISPLAY_ID)).thenReturn(true);
+
+        assertTrue(mPipDesktopState.shouldExitPipExitDesktopMode());
+    }
+
+    @Test
+    public void getOutPipWindowingMode_exitToDesktop_displayFreeform_returnsUndefined() {
+        // Set visible task count to 1 so isPipExitingToDesktopMode returns true
+        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(1);
+        setDisplayWindowingMode(WINDOWING_MODE_FREEFORM);
+
+        assertEquals(WINDOWING_MODE_UNDEFINED, mPipDesktopState.getOutPipWindowingMode());
+    }
+
+    @Test
+    public void getOutPipWindowingMode_exitToDesktop_displayFullscreen_returnsFreeform() {
+        // Set visible task count to 1 so isPipExitingToDesktopMode returns true
+        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(1);
+        setDisplayWindowingMode(WINDOWING_MODE_FULLSCREEN);
+
+        assertEquals(WINDOWING_MODE_FREEFORM, mPipDesktopState.getOutPipWindowingMode());
+    }
+
+    @Test
+    public void getOutPipWindowingMode_exitToFullscreen_displayFullscreen_returnsUndefined() {
+        setDisplayWindowingMode(WINDOWING_MODE_FULLSCREEN);
+
+        assertEquals(WINDOWING_MODE_UNDEFINED, mPipDesktopState.getOutPipWindowingMode());
+    }
+
+    private void setDisplayWindowingMode(int windowingMode) {
+        mDefaultTda.configuration.windowConfiguration.setWindowingMode(windowingMode);
+    }
+}
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 bd857c7dcd45..8e0381e4f933 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
@@ -26,7 +26,6 @@ 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;
@@ -39,12 +38,9 @@ import android.window.WindowContainerTransaction;
 
 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.common.pip.PipDesktopState;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.pip2.PipSurfaceTransactionHelper;
 import com.android.wm.shell.pip2.animation.PipAlphaAnimator;
@@ -55,11 +51,8 @@ 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;
-
 /**
  * Unit test against {@link PipScheduler}
  */
@@ -77,16 +70,13 @@ public class PipSchedulerTest {
     @Mock private PipBoundsState mMockPipBoundsState;
     @Mock private ShellExecutor mMockMainExecutor;
     @Mock private PipTransitionState mMockPipTransitionState;
+    @Mock private PipDesktopState mMockPipDesktopState;
     @Mock private PipTransitionController mMockPipTransitionController;
     @Mock private Runnable mMockUpdateMovementBoundsRunnable;
     @Mock private WindowContainerToken mMockPipTaskToken;
     @Mock private PipSurfaceTransactionHelper.SurfaceControlTransactionFactory mMockFactory;
     @Mock private SurfaceControl.Transaction mMockTransaction;
     @Mock private PipAlphaAnimator mMockAlphaAnimator;
-    @Mock private DesktopUserRepositories mMockDesktopUserRepositories;
-    @Mock private DesktopWallpaperActivityTokenProvider mMockDesktopWallpaperActivityTokenProvider;
-    @Mock private RootTaskDisplayAreaOrganizer mRootTaskDisplayAreaOrganizer;
-
     @Captor private ArgumentCaptor<Runnable> mRunnableArgumentCaptor;
     @Captor private ArgumentCaptor<WindowContainerTransaction> mWctArgumentCaptor;
 
@@ -101,14 +91,9 @@ 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, Optional.of(mMockDesktopUserRepositories),
-                Optional.of(mMockDesktopWallpaperActivityTokenProvider),
-                mRootTaskDisplayAreaOrganizer);
+                mMockPipTransitionState, mMockPipDesktopState);
         mPipScheduler.setPipTransitionController(mMockPipTransitionController);
         mPipScheduler.setSurfaceControlTransactionFactory(mMockFactory);
         mPipScheduler.setPipAlphaAnimatorSupplier((context, leash, startTx, finishTx, direction) ->