Allow split pair to move to external display

This CL allows split pair to move from primary display to external
display, when the flags are enabled.

Bug: 392967255
Test: manual
Flag: com.android.window.flags.enable_non_default_display_split
Recall: http://recall/clips/c1312d26-a9ee-4f10-9ff7-1fc55376365b
Change-Id: Ida41fd34e4522d703b78d17312feded2edeb440d

4 files changed, 62 insertions, 1 deletions

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 10f87058b527..ffec3ec673ec 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
@@ -1015,6 +1015,23 @@ class DesktopTasksController(
         }
 
         val wct = WindowContainerTransaction()
+
+        // check if the task is part of splitscreen
+        if (
+            Flags.enableNonDefaultDisplaySplit() &&
+                Flags.enableMoveToNextDisplayShortcut() &&
+                splitScreenController.isTaskInSplitScreen(task.taskId)
+        ) {
+            val stageCoordinatorRootTaskToken =
+                splitScreenController.multiDisplayProvider.getDisplayRootForDisplayId(
+                    DEFAULT_DISPLAY
+                )
+
+            wct.reparent(stageCoordinatorRootTaskToken, displayAreaInfo.token, true /* onTop */)
+            transitions.startTransition(TRANSIT_CHANGE, wct, /* handler= */ null)
+            return
+        }
+
         if (!task.isFreeform) {
             addMoveToDesktopChanges(wct, task, displayId)
         } else if (Flags.enableMoveToNextDisplayShortcut()) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitMultiDisplayProvider.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitMultiDisplayProvider.java
new file mode 100644
index 000000000000..d2e57e51762b
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitMultiDisplayProvider.java
@@ -0,0 +1,29 @@
+/*
+ * 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.splitscreen;
+
+import android.window.WindowContainerToken;
+
+public interface SplitMultiDisplayProvider {
+    /**
+     * Returns the WindowContainerToken for the root of the given display ID.
+     *
+     * @param displayId The ID of the display.
+     * @return The {@link WindowContainerToken} associated with the display's root task.
+     */
+    WindowContainerToken getDisplayRootForDisplayId(int displayId);
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index ae0159263364..e9f8a4a86d27 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -321,6 +321,10 @@ public class SplitScreenController implements SplitDragPolicy.Starter,
         return mStageCoordinator;
     }
 
+    public SplitMultiDisplayProvider getMultiDisplayProvider() {
+        return mStageCoordinator;
+    }
+
     @Nullable
     public ActivityManager.RunningTaskInfo getTaskInfo(@SplitPosition int splitPosition) {
         if (!isSplitScreenVisible() || splitPosition == SPLIT_POSITION_UNDEFINED) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index 2174017996a8..7b139cb2a695 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -189,7 +189,8 @@ import java.util.function.Predicate;
  */
 public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         DisplayController.OnDisplaysChangedListener, Transitions.TransitionHandler,
-        ShellTaskOrganizer.TaskListener, StageTaskListener.StageListenerCallbacks {
+        ShellTaskOrganizer.TaskListener, StageTaskListener.StageListenerCallbacks,
+        SplitMultiDisplayProvider {
 
     private static final String TAG = StageCoordinator.class.getSimpleName();
 
@@ -287,6 +288,16 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         mSplitTransitions.registerSplitAnimListener(listener, executor);
     }
 
+    @Override
+    public WindowContainerToken getDisplayRootForDisplayId(int displayId) {
+        if (displayId == DEFAULT_DISPLAY) {
+            return mRootTaskInfo != null ? mRootTaskInfo.token : null;
+        }
+
+        // TODO(b/393217881): support different root task on external displays.
+        return null; // Return null for unknown display IDs
+    }
+
     class SplitRequest {
         @SplitPosition
         int mActivatePosition;