diff --git a/core/java/android/app/ActivityThread.java b/core/java/android/app/ActivityThread.java
index 8d39d8a..7149096 100644
--- a/core/java/android/app/ActivityThread.java
+++ b/core/java/android/app/ActivityThread.java
@@ -237,7 +237,6 @@
 import java.util.TimeZone;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 
 /**
@@ -338,6 +337,11 @@
      */
     @UnsupportedAppUsage
     final ArrayMap<IBinder, ActivityClientRecord> mActivities = new ArrayMap<>();
+    /**
+     * Maps from activity token to local record of the activities that are preparing to be launched.
+     */
+    final Map<IBinder, ActivityClientRecord> mLaunchingActivities =
+            Collections.synchronizedMap(new ArrayMap<IBinder, ActivityClientRecord>());
     /** The activities to be truly destroyed (not include relaunch). */
     final Map<IBinder, ClientTransactionItem> mActivitiesToBeDestroyed =
             Collections.synchronizedMap(new ArrayMap<IBinder, ClientTransactionItem>());
@@ -347,7 +351,6 @@
     // Number of activities that are currently visible on-screen.
     @UnsupportedAppUsage
     int mNumVisibleActivities = 0;
-    private final AtomicInteger mNumLaunchingActivities = new AtomicInteger();
     @GuardedBy("mAppThread")
     private int mLastProcessState = PROCESS_STATE_UNKNOWN;
     @GuardedBy("mAppThread")
@@ -3256,6 +3259,21 @@
     }
 
     @Override
+    public void addLaunchingActivity(IBinder token, ActivityClientRecord activity) {
+        mLaunchingActivities.put(token, activity);
+    }
+
+    @Override
+    public ActivityClientRecord getLaunchingActivity(IBinder token) {
+        return mLaunchingActivities.get(token);
+    }
+
+    @Override
+    public void removeLaunchingActivity(IBinder token) {
+        mLaunchingActivities.remove(token);
+    }
+
+    @Override
     public ActivityClientRecord getActivityClient(IBinder token) {
         return mActivities.get(token);
     }
@@ -3299,7 +3317,7 @@
             // Defer the top state for VM to avoid aggressive JIT compilation affecting activity
             // launch time.
             if (processState == ActivityManager.PROCESS_STATE_TOP
-                    && mNumLaunchingActivities.get() > 0) {
+                    && !mLaunchingActivities.isEmpty()) {
                 mPendingProcessState = processState;
                 mH.postDelayed(this::applyPendingProcessState, PENDING_TOP_PROCESS_STATE_TIMEOUT);
             } else {
@@ -3315,7 +3333,7 @@
         // Handle the pending configuration if the process state is changed from cached to
         // non-cached. Except the case where there is a launching activity because the
         // LaunchActivityItem will handle it.
-        if (wasCached && !isCachedProcessState() && mNumLaunchingActivities.get() == 0) {
+        if (wasCached && !isCachedProcessState() && mLaunchingActivities.isEmpty()) {
             final Configuration pendingConfig =
                     mConfigurationController.getPendingConfiguration(false /* clearPending */);
             if (pendingConfig == null) {
@@ -3353,11 +3371,6 @@
         }
     }
 
-    @Override
-    public void countLaunchingActivities(int num) {
-        mNumLaunchingActivities.getAndAdd(num);
-    }
-
     @UnsupportedAppUsage
     public final void sendActivityResult(
             IBinder token, String id, int requestCode,
diff --git a/core/java/android/app/ClientTransactionHandler.java b/core/java/android/app/ClientTransactionHandler.java
index c752f34..115101c 100644
--- a/core/java/android/app/ClientTransactionHandler.java
+++ b/core/java/android/app/ClientTransactionHandler.java
@@ -82,9 +82,6 @@
     /** Set current process state. */
     public abstract void updateProcessState(int processState, boolean fromIpc);
 
-    /** Count how many activities are launching. */
-    public abstract void countLaunchingActivities(int num);
-
     // Execute phase related logic and handlers. Methods here execute actual lifecycle transactions
     // and deliver callbacks.
 
@@ -193,6 +190,26 @@
             FixedRotationAdjustments fixedRotationAdjustments);
 
     /**
+     * Add {@link ActivityClientRecord} that is preparing to be launched.
+     * @param token Activity token.
+     * @param activity An initialized instance of {@link ActivityClientRecord} to use during launch.
+     */
+    public abstract void addLaunchingActivity(IBinder token, ActivityClientRecord activity);
+
+    /**
+     * Get {@link ActivityClientRecord} that is preparing to be launched.
+     * @param token Activity token.
+     * @return An initialized instance of {@link ActivityClientRecord} to use during launch.
+     */
+    public abstract ActivityClientRecord getLaunchingActivity(IBinder token);
+
+    /**
+     * Remove {@link ActivityClientRecord} from the launching activity list.
+     * @param token Activity token.
+     */
+    public abstract void removeLaunchingActivity(IBinder token);
+
+    /**
      * Get {@link android.app.ActivityThread.ActivityClientRecord} instance that corresponds to the
      * provided token.
      */
diff --git a/core/java/android/app/servertransaction/ActivityConfigurationChangeItem.java b/core/java/android/app/servertransaction/ActivityConfigurationChangeItem.java
index d5585db..032b57e 100644
--- a/core/java/android/app/servertransaction/ActivityConfigurationChangeItem.java
+++ b/core/java/android/app/servertransaction/ActivityConfigurationChangeItem.java
@@ -40,7 +40,8 @@
 
     @Override
     public void preExecute(android.app.ClientTransactionHandler client, IBinder token) {
-        final ActivityClientRecord r = getActivityClientRecord(client, token);
+        final ActivityClientRecord r = getActivityClientRecord(client, token,
+                true /* includeLaunching */);
         // Notify the client of an upcoming change in the token configuration. This ensures that
         // batches of config change items only process the newest configuration.
         client.updatePendingActivityConfiguration(r, mConfiguration);
diff --git a/core/java/android/app/servertransaction/ActivityTransactionItem.java b/core/java/android/app/servertransaction/ActivityTransactionItem.java
index f7d7e9d..a539812 100644
--- a/core/java/android/app/servertransaction/ActivityTransactionItem.java
+++ b/core/java/android/app/servertransaction/ActivityTransactionItem.java
@@ -55,15 +55,40 @@
 
     @NonNull ActivityClientRecord getActivityClientRecord(
             @NonNull ClientTransactionHandler client, IBinder token) {
-        final ActivityClientRecord r = client.getActivityClient(token);
+        return getActivityClientRecord(client, token, false /* includeLaunching */);
+    }
+
+    /**
+     * Get the {@link ActivityClientRecord} instance that corresponds to the provided token.
+     * @param client Target client handler.
+     * @param token Target activity token.
+     * @param includeLaunching Indicate to also find the {@link ActivityClientRecord} in launching
+     *                         activity list. It should be noted that there is no activity in
+     *                         {@link ActivityClientRecord} from the launching activity list.
+     * @return The {@link ActivityClientRecord} instance that corresponds to the provided token.
+     */
+    @NonNull ActivityClientRecord getActivityClientRecord(
+            @NonNull ClientTransactionHandler client, IBinder token, boolean includeLaunching) {
+        ActivityClientRecord r = client.getActivityClient(token);
+        if (r != null) {
+            if (client.getActivity(token) == null) {
+                throw new IllegalArgumentException("Activity must not be null to execute "
+                        + "transaction item");
+            }
+            return r;
+        }
+        // The activity may not be launched yet. Fallback to check launching activity.
+        if (includeLaunching) {
+            r = client.getLaunchingActivity(token);
+        }
         if (r == null) {
             throw new IllegalArgumentException("Activity client record must not be null to execute "
                     + "transaction item");
         }
-        if (client.getActivity(token) == null) {
-            throw new IllegalArgumentException("Activity must not be null to execute "
-                    + "transaction item");
-        }
+
+        // We don't need to check the activity of launching activity client records because they
+        // have not been launched yet.
+
         return r;
     }
 }
diff --git a/core/java/android/app/servertransaction/LaunchActivityItem.java b/core/java/android/app/servertransaction/LaunchActivityItem.java
index e281a02..34e4fcd 100644
--- a/core/java/android/app/servertransaction/LaunchActivityItem.java
+++ b/core/java/android/app/servertransaction/LaunchActivityItem.java
@@ -82,7 +82,12 @@
 
     @Override
     public void preExecute(ClientTransactionHandler client, IBinder token) {
-        client.countLaunchingActivities(1);
+        ActivityClientRecord r = new ActivityClientRecord(token, mIntent, mIdent, mInfo,
+                mOverrideConfig, mCompatInfo, mReferrer, mVoiceInteractor, mState, mPersistentState,
+                mPendingResults, mPendingNewIntents, mActivityOptions, mIsForward, mProfilerInfo,
+                client, mAssistToken, mFixedRotationAdjustments, mShareableActivityToken,
+                mLaunchedFromBubble);
+        client.addLaunchingActivity(token, r);
         client.updateProcessState(mProcState, false);
         client.updatePendingConfiguration(mCurConfig);
         if (mActivityClientController != null) {
@@ -94,11 +99,7 @@
     public void execute(ClientTransactionHandler client, IBinder token,
             PendingTransactionActions pendingActions) {
         Trace.traceBegin(TRACE_TAG_ACTIVITY_MANAGER, "activityStart");
-        ActivityClientRecord r = new ActivityClientRecord(token, mIntent, mIdent, mInfo,
-                mOverrideConfig, mCompatInfo, mReferrer, mVoiceInteractor, mState, mPersistentState,
-                mPendingResults, mPendingNewIntents, mActivityOptions, mIsForward, mProfilerInfo,
-                client, mAssistToken, mFixedRotationAdjustments, mShareableActivityToken,
-                mLaunchedFromBubble);
+        ActivityClientRecord r = client.getLaunchingActivity(token);
         client.handleLaunchActivity(r, pendingActions, null /* customIntent */);
         Trace.traceEnd(TRACE_TAG_ACTIVITY_MANAGER);
     }
@@ -106,7 +107,7 @@
     @Override
     public void postExecute(ClientTransactionHandler client, IBinder token,
             PendingTransactionActions pendingActions) {
-        client.countLaunchingActivities(-1);
+        client.removeLaunchingActivity(token);
     }
 
 
diff --git a/core/java/android/app/servertransaction/MoveToDisplayItem.java b/core/java/android/app/servertransaction/MoveToDisplayItem.java
index 944367e..4b8a347 100644
--- a/core/java/android/app/servertransaction/MoveToDisplayItem.java
+++ b/core/java/android/app/servertransaction/MoveToDisplayItem.java
@@ -40,7 +40,8 @@
 
     @Override
     public void preExecute(ClientTransactionHandler client, IBinder token) {
-        final ActivityClientRecord r = getActivityClientRecord(client, token);
+        final ActivityClientRecord r = getActivityClientRecord(client, token,
+                true /* includeLaunching */);
         // Notify the client of an upcoming change in the token configuration. This ensures that
         // batches of config change items only process the newest configuration.
         client.updatePendingActivityConfiguration(r, mConfiguration);