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 746bfad56ea8e..f102f8e2f5dff 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
@@ -936,7 +936,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                         new IRemoteAnimationFinishedCallback.Stub() {
                             @Override
                             public void onAnimationFinished() throws RemoteException {
-                                onRemoteAnimationFinishedOrCancelled(false /* cancel */, evictWct);
+                                onRemoteAnimationFinishedOrCancelled(evictWct);
                                 finishedCallback.onAnimationFinished();
                             }
                         };
@@ -952,7 +952,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
 
             @Override
             public void onAnimationCancelled(boolean isKeyguardOccluded) {
-                onRemoteAnimationFinishedOrCancelled(true /* cancel */, evictWct);
+                onRemoteAnimationFinishedOrCancelled(evictWct);
                 try {
                     adapter.getRunner().onAnimationCancelled(isKeyguardOccluded);
                 } catch (RemoteException e) {
@@ -973,15 +973,14 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         }
     }
 
-    private void onRemoteAnimationFinishedOrCancelled(boolean cancel,
-            WindowContainerTransaction evictWct) {
+    private void onRemoteAnimationFinishedOrCancelled(WindowContainerTransaction evictWct) {
         mIsSplitEntering = false;
         mShouldUpdateRecents = true;
         mSplitRequest = null;
         // If any stage has no child after animation finished, it means that split will display
         // nothing, such status will happen if task and intent is same app but not support
         // multi-instance, we should exit split and expand that app as full screen.
-        if (!cancel && (mMainStage.getChildCount() == 0 || mSideStage.getChildCount() == 0)) {
+        if (mMainStage.getChildCount() == 0 || mSideStage.getChildCount() == 0) {
             mMainExecutor.execute(() ->
                     exitSplitScreen(mMainStage.getChildCount() == 0
                             ? mSideStage : mMainStage, EXIT_REASON_UNKNOWN));
@@ -1236,8 +1235,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             // Notify recents if we are exiting in a way that breaks the pair, and disable further
             // updates to splits in the recents until we enter split again
             if (shouldBreakPairedTaskInRecents(exitReason) && mShouldUpdateRecents) {
-                recentTasks.removeSplitPair(mMainStage.getLastVisibleTaskId());
-                recentTasks.removeSplitPair(mSideStage.getLastVisibleTaskId());
+                recentTasks.removeSplitPair(mMainStage.getTopVisibleChildTaskId());
+                recentTasks.removeSplitPair(mSideStage.getTopVisibleChildTaskId());
             }
         });
         mShouldUpdateRecents = false;
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 0359761388dc2..a841b7f96d3c5 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
@@ -92,7 +92,6 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
     protected SurfaceControl mDimLayer;
     protected SparseArray<ActivityManager.RunningTaskInfo> mChildrenTaskInfo = new SparseArray<>();
     private final SparseArray<SurfaceControl> mChildrenLeashes = new SparseArray<>();
-    private int mLastVisibleTaskId = INVALID_TASK_ID;
     // TODO(b/204308910): Extracts SplitDecorManager related code to common package.
     private SplitDecorManager mSplitDecorManager;
 
@@ -123,13 +122,6 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
         return contains(t -> t.token.asBinder() == binder);
     }
 
-    /**
-     * Returns the last visible task's id.
-     */
-    int getLastVisibleTaskId() {
-        return mLastVisibleTaskId;
-    }
-
     /**
      * Returns the top visible child task's id.
      */
@@ -229,9 +221,6 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
                 return;
             }
             mChildrenTaskInfo.put(taskInfo.taskId, taskInfo);
-            if (taskInfo.isVisible && taskInfo.taskId != mLastVisibleTaskId) {
-                mLastVisibleTaskId = taskInfo.taskId;
-            }
             mCallbacks.onChildTaskStatusChanged(taskInfo.taskId, true /* present */,
                     taskInfo.isVisible);
             if (!ENABLE_SHELL_TRANSITIONS) {
@@ -264,9 +253,6 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
         } else if (mChildrenTaskInfo.contains(taskId)) {
             mChildrenTaskInfo.remove(taskId);
             mChildrenLeashes.remove(taskId);
-            if (taskId == mLastVisibleTaskId) {
-                mLastVisibleTaskId = INVALID_TASK_ID;
-            }
             mCallbacks.onChildTaskStatusChanged(taskId, false /* present */, taskInfo.isVisible);
             if (ENABLE_SHELL_TRANSITIONS) {
                 // Status is managed/synchronized by the transition lifecycle.