diff --git a/services/core/java/com/android/server/wm/TaskDisplayArea.java b/services/core/java/com/android/server/wm/TaskDisplayArea.java
index cda8c4b78b0cd..87f685e95fe13 100644
--- a/services/core/java/com/android/server/wm/TaskDisplayArea.java
+++ b/services/core/java/com/android/server/wm/TaskDisplayArea.java
@@ -1274,7 +1274,15 @@ final class TaskDisplayArea extends DisplayArea<WindowContainer> {
 
         for (int i = mLaunchRootTasks.size() - 1; i >= 0; --i) {
             if (mLaunchRootTasks.get(i).contains(windowingMode, activityType)) {
-                return mLaunchRootTasks.get(i).task;
+                final Task launchRootTask = mLaunchRootTasks.get(i).task;
+                // Return the focusable root task for improving the UX with staged split screen.
+                final Task adjacentRootTask = launchRootTask != null
+                        ? launchRootTask.mAdjacentTask : null;
+                if (adjacentRootTask != null && adjacentRootTask.isFocusedRootTaskOnDisplay()) {
+                    return adjacentRootTask;
+                } else {
+                    return launchRootTask;
+                }
             }
         }
         return null;
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskDisplayAreaTests.java b/services/tests/wmtests/src/com/android/server/wm/TaskDisplayAreaTests.java
index 9289ce41cd1e4..67b273a5a82d1 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskDisplayAreaTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskDisplayAreaTests.java
@@ -100,6 +100,29 @@ public class TaskDisplayAreaTests extends WindowTestsBase {
         assertNull(actualRootTask);
     }
 
+    @Test
+    public void getLaunchRootTask_checksFocusedRootTask() {
+        final TaskDisplayArea taskDisplayArea = mRootWindowContainer.getDefaultTaskDisplayArea();
+        final Task rootTask = createTaskWithActivity(
+                taskDisplayArea,
+                WINDOWING_MODE_MULTI_WINDOW, ACTIVITY_TYPE_STANDARD, ON_TOP, true);
+        rootTask.mCreatedByOrganizer = true;
+
+        final Task adjacentRootTask = createTask(
+                mDisplayContent, WINDOWING_MODE_MULTI_WINDOW, ACTIVITY_TYPE_STANDARD);
+        adjacentRootTask.mCreatedByOrganizer = true;
+        adjacentRootTask.mAdjacentTask = rootTask;
+        rootTask.mAdjacentTask = adjacentRootTask;
+
+        taskDisplayArea.setLaunchRootTask(rootTask,
+                new int[]{WINDOWING_MODE_MULTI_WINDOW}, new int[]{ACTIVITY_TYPE_STANDARD});
+
+        Task actualRootTask = taskDisplayArea.getLaunchRootTask(
+                WINDOWING_MODE_MULTI_WINDOW, ACTIVITY_TYPE_STANDARD, null /* options */,
+                null /* sourceTask */, 0 /*launchFlags*/);
+        assertTrue(actualRootTask.isFocusedRootTaskOnDisplay());
+    }
+
     @Test
     public void getLaunchRootTask_fromLaunchAdjacentFlagRoot_checksAdjacentRoot() {
         final ActivityRecord activity = createNonAttachedActivityRecord(mDisplayContent);