diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java
deleted file mode 100644
index 1dad413..0000000
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java
+++ /dev/null
@@ -1,85 +0,0 @@
-/*
- * Copyright (C) 2020 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 static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN;
-
-import android.content.Context;
-import android.view.SurfaceSession;
-import android.window.WindowContainerToken;
-import android.window.WindowContainerTransaction;
-
-import com.android.internal.protolog.ProtoLog;
-import com.android.launcher3.icons.IconProvider;
-import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.common.SyncTransactionQueue;
-import com.android.wm.shell.windowdecor.WindowDecorViewModel;
-
-import java.util.Optional;
-
-/**
- * Main stage for split-screen mode. When split-screen is active all standard activity types launch
- * on the main stage, except for task that are explicitly pinned to the {@link StageTaskListener}.
- * @see StageCoordinator
- */
-class MainStage extends StageTaskListener {
-    private boolean mIsActive = false;
-
-    MainStage(Context context, ShellTaskOrganizer taskOrganizer, int displayId,
-            StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue,
-            SurfaceSession surfaceSession, IconProvider iconProvider,
-            Optional<WindowDecorViewModel> windowDecorViewModel) {
-        super(context, taskOrganizer, displayId, callbacks, syncQueue, surfaceSession,
-                iconProvider, windowDecorViewModel);
-    }
-
-    boolean isActive() {
-        return mIsActive;
-    }
-
-    void activate(WindowContainerTransaction wct, boolean includingTopTask) {
-        if (mIsActive) return;
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "activate: main stage includingTopTask=%b",
-                includingTopTask);
-
-        if (includingTopTask) {
-            reparentTopTask(wct);
-        }
-
-        mIsActive = true;
-    }
-
-    void deactivate(WindowContainerTransaction wct) {
-        deactivate(wct, false /* toTop */);
-    }
-
-    void deactivate(WindowContainerTransaction wct, boolean toTop) {
-        if (!mIsActive) return;
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "deactivate: main stage toTop=%b rootTaskInfo=%s",
-                toTop, mRootTaskInfo);
-        mIsActive = false;
-
-        if (mRootTaskInfo == null) return;
-        final WindowContainerToken rootToken = mRootTaskInfo.token;
-        wct.reparentTasks(
-                rootToken,
-                null /* newParent */,
-                null /* windowingModes */,
-                null /* activityTypes */,
-                toTop);
-    }
-}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
index 526c1d4..b36b1f8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
@@ -44,7 +44,7 @@
     int STAGE_TYPE_UNDEFINED = -1;
     /**
      * The main stage type.
-     * @see MainStage
+     * @see StageTaskListener
      */
     int STAGE_TYPE_MAIN = 0;
 
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 f3959cc..f3113dc 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
@@ -34,6 +34,7 @@
 import static android.window.TransitionInfo.FLAG_IS_DISPLAY;
 import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER;
 
+import static com.android.wm.shell.Flags.enableFlexibleSplit;
 import static com.android.wm.shell.common.split.SplitLayout.PARALLAX_ALIGN_CENTER;
 import static com.android.wm.shell.common.split.SplitScreenUtils.reverseSplitPosition;
 import static com.android.wm.shell.common.split.SplitScreenUtils.splitFailureMessage;
@@ -154,14 +155,12 @@
 import java.util.concurrent.Executor;
 
 /**
- * Coordinates the staging (visibility, sizing, ...) of the split-screen {@link MainStage} and
- * other stages.
+ * Coordinates the staging (visibility, sizing, ...) of the split-screen stages.
  * Some high-level rules:
  * - The {@link StageCoordinator} is only considered active if the other stages contain at
  * least one child task.
- * - The {@link MainStage} should only have children if the coordinator is active.
- * - The {@link SplitLayout} divider is only visible if both the {@link MainStage}
- * and other stages are visible.
+ * - The {@link SplitLayout} divider is only visible if multiple {@link StageTaskListener}s are
+ * visible
  * - Both stages are put under a single-top root task.
  * This rules are mostly implemented in {@link #onStageVisibilityChanged(StageListenerImpl)} and
  * {@link #onStageHasChildrenChanged(StageListenerImpl).}
@@ -174,7 +173,7 @@
 
     private final SurfaceSession mSurfaceSession = new SurfaceSession();
 
-    private final MainStage mMainStage;
+    private final StageTaskListener mMainStage;
     private final StageListenerImpl mMainStageListener = new StageListenerImpl();
     private final StageTaskListener mSideStage;
     private final StageListenerImpl mSideStageListener = new StageListenerImpl();
@@ -329,7 +328,7 @@
         taskOrganizer.createRootTask(displayId, WINDOWING_MODE_FULLSCREEN, this /* listener */);
 
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "Creating main/side root task");
-        mMainStage = new MainStage(
+        mMainStage = new StageTaskListener(
                 mContext,
                 mTaskOrganizer,
                 mDisplayId,
@@ -367,8 +366,9 @@
 
     @VisibleForTesting
     StageCoordinator(Context context, int displayId, SyncTransactionQueue syncQueue,
-            ShellTaskOrganizer taskOrganizer, MainStage mainStage, StageTaskListener sideStage,
-            DisplayController displayController, DisplayImeController displayImeController,
+            ShellTaskOrganizer taskOrganizer, StageTaskListener mainStage,
+            StageTaskListener sideStage, DisplayController displayController,
+            DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController, SplitLayout splitLayout,
             Transitions transitions, TransactionPool transactionPool, ShellExecutor mainExecutor,
             Handler mainHandler, Optional<RecentTasksController> recentTasks,
@@ -420,6 +420,10 @@
         return mSideStageListener.mVisible && mMainStageListener.mVisible;
     }
 
+    private void activateSplit(WindowContainerTransaction wct, boolean includingTopTask) {
+        mMainStage.activate(wct, includingTopTask);
+    }
+
     public boolean isSplitActive() {
         return mMainStage.isActive();
     }
@@ -505,10 +509,10 @@
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "removeFromSideStage: task=%d", taskId);
         final WindowContainerTransaction wct = new WindowContainerTransaction();
 
-        /**
-         * {@link MainStage} will be deactivated in {@link #onStageHasChildrenChanged} if the
-         * other stages no longer have children.
-         */
+
+        // MainStage will be deactivated in onStageHasChildrenChanged() if the other stages
+        // no longer have children.
+
         final boolean result = mSideStage.removeTask(taskId,
                 isSplitActive() ? mMainStage.mRootTaskInfo.token : null,
                 wct);
@@ -805,7 +809,7 @@
         if (!isSplitActive()) {
             // Build a request WCT that will launch both apps such that task 0 is on the main stage
             // while task 1 is on the side stage.
-            mMainStage.activate(wct, false /* reparent */);
+            activateSplit(wct, false /* reparentToTop */);
         }
         mSplitLayout.setDivideRatio(snapPosition);
         updateWindowBounds(mSplitLayout, wct);
@@ -872,7 +876,7 @@
         if (!isSplitActive()) {
             // Build a request WCT that will launch both apps such that task 0 is on the main stage
             // while task 1 is on the side stage.
-            mMainStage.activate(wct, false /* reparent */);
+            activateSplit(wct, false /* reparentToTop */);
         }
 
         setSideStagePosition(splitPosition, wct);
@@ -1439,7 +1443,7 @@
             setSideStagePosition(startPosition, wct);
             mSideStage.addTask(taskInfo, wct);
         }
-        mMainStage.activate(wct, true /* includingTopTask */);
+        activateSplit(wct, true /* reparentToTop */);
         prepareSplitLayout(wct, resizeAnim);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
index 4593553..29b5114 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
@@ -72,6 +72,10 @@
 public class StageTaskListener implements ShellTaskOrganizer.TaskListener {
     private static final String TAG = StageTaskListener.class.getSimpleName();
 
+    // No current way to enforce this but if enableFlexibleSplit() is enabled, then only 1 of the
+    // stages should have this be set/being used
+    private boolean mIsActive;
+
     /** Callback interface for listening to changes in a split-screen stage. */
     public interface StageListenerCallbacks {
         void onRootTaskAppeared();
@@ -475,6 +479,44 @@
         });
     }
 
+    // ---------
+    // Previously only used in MainStage
+    boolean isActive() {
+        return mIsActive;
+    }
+
+    void activate(WindowContainerTransaction wct, boolean includingTopTask) {
+        if (mIsActive) return;
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "activate: includingTopTask=%b",
+                includingTopTask);
+
+        if (includingTopTask) {
+            reparentTopTask(wct);
+        }
+
+        mIsActive = true;
+    }
+
+    void deactivate(WindowContainerTransaction wct) {
+        deactivate(wct, false /* toTop */);
+    }
+
+    void deactivate(WindowContainerTransaction wct, boolean toTop) {
+        if (!mIsActive) return;
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "deactivate: toTop=%b rootTaskInfo=%s",
+                toTop, mRootTaskInfo);
+        mIsActive = false;
+
+        if (mRootTaskInfo == null) return;
+        final WindowContainerToken rootToken = mRootTaskInfo.token;
+        wct.reparentTasks(
+                rootToken,
+                null /* newParent */,
+                null /* windowingModes */,
+                null /* activityTypes */,
+                toTop);
+    }
+
     // --------
     // Previously only used in SideStage
     boolean removeAllTasks(WindowContainerTransaction wct, boolean toTop) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java
deleted file mode 100644
index b1befc4..0000000
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java
+++ /dev/null
@@ -1,78 +0,0 @@
-/*
- * Copyright (C) 2020 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 static android.view.Display.DEFAULT_DISPLAY;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import android.app.ActivityManager;
-import android.view.SurfaceControl;
-import android.view.SurfaceSession;
-import android.window.WindowContainerTransaction;
-
-import androidx.test.annotation.UiThreadTest;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
-import com.android.launcher3.icons.IconProvider;
-import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
-import com.android.wm.shell.TestRunningTaskInfoBuilder;
-import com.android.wm.shell.common.SyncTransactionQueue;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-
-import java.util.Optional;
-
-/** Tests for {@link MainStage} */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
-public class MainStageTests extends ShellTestCase {
-    @Mock private ShellTaskOrganizer mTaskOrganizer;
-    @Mock private StageTaskListener.StageListenerCallbacks mCallbacks;
-    @Mock private SyncTransactionQueue mSyncQueue;
-    @Mock private ActivityManager.RunningTaskInfo mRootTaskInfo;
-    @Mock private SurfaceControl mRootLeash;
-    @Mock private IconProvider mIconProvider;
-    private WindowContainerTransaction mWct = new WindowContainerTransaction();
-    private SurfaceSession mSurfaceSession = new SurfaceSession();
-    private MainStage mMainStage;
-
-    @Before
-    @UiThreadTest
-    public void setup() {
-        MockitoAnnotations.initMocks(this);
-        mRootTaskInfo = new TestRunningTaskInfoBuilder().build();
-        mMainStage = new MainStage(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mCallbacks,
-                mSyncQueue, mSurfaceSession, mIconProvider, Optional.empty());
-        mMainStage.onTaskAppeared(mRootTaskInfo, mRootLeash);
-    }
-
-    @Test
-    public void testActiveDeactivate() {
-        mMainStage.activate(mWct, true /* reparent */);
-        assertThat(mMainStage.isActive()).isTrue();
-
-        mMainStage.deactivate(mWct);
-        assertThat(mMainStage.isActive()).isFalse();
-    }
-}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
index 4de2278..751275b 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
@@ -74,10 +74,10 @@
         final SurfaceControl mRootLeash;
 
         TestStageCoordinator(Context context, int displayId, SyncTransactionQueue syncQueue,
-                ShellTaskOrganizer taskOrganizer, MainStage mainStage, StageTaskListener sideStage,
-                DisplayController displayController, DisplayImeController imeController,
-                DisplayInsetsController insetsController, SplitLayout splitLayout,
-                Transitions transitions, TransactionPool transactionPool,
+                ShellTaskOrganizer taskOrganizer, StageTaskListener mainStage,
+                StageTaskListener sideStage, DisplayController displayController,
+                DisplayImeController imeController, DisplayInsetsController insetsController,
+                SplitLayout splitLayout, Transitions transitions, TransactionPool transactionPool,
                 ShellExecutor mainExecutor, Handler mainHandler,
                 Optional<RecentTasksController> recentTasks,
                 LaunchAdjacentController launchAdjacentController,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
index e167433..af288c8 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
@@ -116,7 +116,7 @@
     @Mock private SplitScreen.SplitInvocationListener mInvocationListener;
     private final TestShellExecutor mTestShellExecutor = new TestShellExecutor();
     private SplitLayout mSplitLayout;
-    private MainStage mMainStage;
+    private StageTaskListener mMainStage;
     private StageTaskListener mSideStage;
     private StageCoordinator mStageCoordinator;
     private SplitScreenTransitions mSplitScreenTransitions;
@@ -133,7 +133,7 @@
         doReturn(mockExecutor).when(mTransitions).getAnimExecutor();
         doReturn(mock(SurfaceControl.Transaction.class)).when(mTransactionPool).acquire();
         mSplitLayout = SplitTestUtils.createMockSplitLayout();
-        mMainStage = spy(new MainStage(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mock(
+        mMainStage = spy(new StageTaskListener(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mock(
                 StageTaskListener.StageListenerCallbacks.class), mSyncQueue, mSurfaceSession,
                 mIconProvider, Optional.of(mWindowDecorViewModel)));
         mMainStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
index c9e1414..ce343b8 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
@@ -97,7 +97,7 @@
     @Mock
     private SyncTransactionQueue mSyncQueue;
     @Mock
-    private MainStage mMainStage;
+    private StageTaskListener mMainStage;
     @Mock
     private StageTaskListener mSideStage;
     @Mock
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
index acd612e..8b5cb97 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
@@ -199,4 +199,13 @@
         assertThat(mStageTaskListener.removeTask(task.taskId, null, mWct)).isTrue();
         verify(mWct).reparent(eq(task.token), isNull(), eq(false));
     }
+
+    @Test
+    public void testActiveDeactivate() {
+        mStageTaskListener.activate(mWct, true /* reparent */);
+        assertThat(mStageTaskListener.isActive()).isTrue();
+
+        mStageTaskListener.deactivate(mWct);
+        assertThat(mStageTaskListener.isActive()).isFalse();
+    }
 }