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 67223c3..5bfecf2 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
@@ -67,6 +67,8 @@
 import com.android.wm.shell.transition.Transitions;
 
 import java.io.PrintWriter;
+import java.util.Arrays;
+import java.util.Comparator;
 import java.util.concurrent.Executor;
 
 /**
@@ -295,11 +297,15 @@
         mRootTDAOrganizer.attachToDisplayArea(DEFAULT_DISPLAY, builder);
         SurfaceControl sc = builder.build();
         SurfaceControl.Transaction transaction = new SurfaceControl.Transaction();
+
+        // Ensure that we order these in the parent in the right z-order as their previous order
+        Arrays.sort(apps, (a1, a2) -> a1.prefixOrderIndex - a2.prefixOrderIndex);
+        int layer = 1;
         for (RemoteAnimationTarget appTarget : apps) {
-            // TODO(b/195958376) set the correct layer/z-order in transaction for the new surface
             transaction.reparent(appTarget.leash, sc);
             transaction.setPosition(appTarget.leash, appTarget.screenSpaceBounds.left,
                     appTarget.screenSpaceBounds.top);
+            transaction.setLayer(appTarget.leash, layer++);
         }
         transaction.apply();
         transaction.close();
diff --git a/services/core/java/com/android/server/wm/InputMonitor.java b/services/core/java/com/android/server/wm/InputMonitor.java
index 9f24b50..f5bfef9 100644
--- a/services/core/java/com/android/server/wm/InputMonitor.java
+++ b/services/core/java/com/android/server/wm/InputMonitor.java
@@ -582,7 +582,7 @@
                 if (recentsAnimationController.updateInputConsumerForApp(
                         mRecentsAnimationInputConsumer.mWindowHandle)) {
                     mRecentsAnimationInputConsumer.show(mInputTransaction,
-                            recentsAnimationController.getHighestLayerActivity());
+                            recentsAnimationController.getHighestLayerWindow());
                     mAddRecentsAnimationInputConsumerHandle = false;
                 }
             }
diff --git a/services/core/java/com/android/server/wm/RecentsAnimationController.java b/services/core/java/com/android/server/wm/RecentsAnimationController.java
index d7dc306..e8d9c15b 100644
--- a/services/core/java/com/android/server/wm/RecentsAnimationController.java
+++ b/services/core/java/com/android/server/wm/RecentsAnimationController.java
@@ -1103,9 +1103,9 @@
     }
 
     /**
-     * Returns the activity with the highest layer, or null if none is found.
+     * Returns the window with the highest layer, or null if none is found.
      */
-    public ActivityRecord getHighestLayerActivity() {
+    public WindowState getHighestLayerWindow() {
         int highestLayer = Integer.MIN_VALUE;
         Task highestLayerTask = null;
         for (int i = mPendingAnimations.size() - 1; i >= 0; i--) {
@@ -1116,7 +1116,7 @@
                 highestLayerTask = adapter.mTask;
             }
         }
-        return highestLayerTask.getTopMostActivity();
+        return highestLayerTask.getTopMostActivity().getTopChild();
     }
 
     boolean isAnimatingTask(Task task) {