diff --git a/services/core/java/com/android/server/wm/ActivityStarter.java b/services/core/java/com/android/server/wm/ActivityStarter.java
index a9192c4c6139e8155d6b5ced36a1b1d43a81605a..d7a696f47b7ea222569d673455fcb629ce96766a 100644
--- a/services/core/java/com/android/server/wm/ActivityStarter.java
+++ b/services/core/java/com/android/server/wm/ActivityStarter.java
@@ -1365,18 +1365,6 @@ class ActivityStarter {
             request.outActivity[0] = mLastStartActivityRecord;
         }
 
-        // Reset the ActivityRecord#mCurrentLaunchCanTurnScreenOn state of activity started
-        // before this one if it is no longer the top-most focusable activity.
-        // Doing so in case the state is not yet consumed during rapid activity launch.
-        if (previousStart != null && !previousStart.finishing && previousStart.isAttached()
-                && previousStart.currentLaunchCanTurnScreenOn()) {
-            final ActivityRecord topFocusable = previousStart.getDisplayContent().getActivity(
-                    ar -> ar.isFocusable() && !ar.finishing);
-            if (previousStart != topFocusable) {
-                previousStart.setCurrentLaunchCanTurnScreenOn(false);
-            }
-        }
-
         return mLastStartActivityResult;
     }
 
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index 7b1c66132c7e2e9a4ff749d51ff6d47b51e45aea..4aa3e3644daa69683dc617900f17bf5b76fa61ad 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -1374,13 +1374,22 @@ class Transition implements BLASTSyncEngine.TransactionReadyListener {
         // processed all the participants first (in particular, we want to trigger pip-enter first)
         for (int i = 0; i < mParticipants.size(); ++i) {
             final ActivityRecord ar = mParticipants.valueAt(i).asActivityRecord();
+            if (ar == null) continue;
+
             // If the activity was just inserted to an invisible task, it will keep INITIALIZING
             // state. Then no need to notify the callback to avoid clearing some states
             // unexpectedly, e.g. launch-task-behind.
-            if (ar != null && (ar.isVisibleRequested()
-                    || !ar.isState(ActivityRecord.State.INITIALIZING))) {
+            if (ar.isVisibleRequested() || !ar.isState(ActivityRecord.State.INITIALIZING)) {
                 mController.dispatchLegacyAppTransitionFinished(ar);
             }
+
+            // Reset the ActivityRecord#mCurrentLaunchCanTurnScreenOn state if it is not the top
+            // running activity. Doing so in case the state is not yet consumed during rapid
+            // activity launch.
+            if (ar.currentLaunchCanTurnScreenOn() && ar.getDisplayContent() != null
+                    && ar.getDisplayContent().topRunningActivity() != ar) {
+                ar.setCurrentLaunchCanTurnScreenOn(false);
+            }
         }
 
         // Update the input-sink (touch-blocking) state now that the animation is finished.