diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index 80d6137..7065f37 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -23,6 +23,8 @@
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT;
+import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_TYPE_A;
+import static com.android.wm.shell.Flags.enableFlexibleSplit;
 import static com.android.wm.shell.shared.GroupedTaskInfo.TYPE_SPLIT;
 
 import android.app.ActivityManager.RunningTaskInfo;
@@ -228,7 +230,8 @@
             return;
         }
 
-        if (stage == SplitConfigurationOptions.STAGE_TYPE_MAIN) {
+        if (stage == SplitConfigurationOptions.STAGE_TYPE_MAIN
+                || (enableFlexibleSplit() && stage == STAGE_TYPE_A)) {
             mMainStagePosition.taskId = taskId;
         } else {
             mSideStagePosition.taskId = taskId;
diff --git a/src/com/android/launcher3/util/SplitConfigurationOptions.java b/src/com/android/launcher3/util/SplitConfigurationOptions.java
index f457e4e..44a7c6f 100644
--- a/src/com/android/launcher3/util/SplitConfigurationOptions.java
+++ b/src/com/android/launcher3/util/SplitConfigurationOptions.java
@@ -73,7 +73,27 @@
      */
     public static final int STAGE_TYPE_SIDE = 1;
 
-    @IntDef({STAGE_TYPE_UNDEFINED, STAGE_TYPE_MAIN, STAGE_TYPE_SIDE})
+    /**
+     * Position independent stage identifier for a given Stage
+     */
+    public static final int STAGE_TYPE_A = 2;
+    /**
+     * Position independent stage identifier for a given Stage
+     */
+    public static final int STAGE_TYPE_B = 3;
+    /**
+     * Position independent stage identifier for a given Stage
+     */
+    public static final int STAGE_TYPE_C = 4;
+
+    @IntDef({
+            STAGE_TYPE_UNDEFINED,
+            STAGE_TYPE_MAIN,
+            STAGE_TYPE_SIDE,
+            STAGE_TYPE_A,
+            STAGE_TYPE_B,
+            STAGE_TYPE_C
+    })
     public @interface StageType {}
     ///////////////////////////////////