diff --git a/services/core/java/com/android/server/pm/PackageInstallerSession.java b/services/core/java/com/android/server/pm/PackageInstallerSession.java
index 020c23d902be9..38059772cfd1a 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerSession.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerSession.java
@@ -2406,13 +2406,6 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
         final VerificationParams verifyingSession = prepareForVerification();
         if (isMultiPackage()) {
             final List<PackageInstallerSession> childSessions = getChildSessions();
-            // Spot check to reject a non-staged multi package install of APEXes and APKs.
-            if (!params.isStaged && containsApkSession()
-                    && sessionContains(s -> s.isApexSession())) {
-                throw new PackageManagerException(
-                    PackageManager.INSTALL_FAILED_SESSION_INVALID,
-                    "Non-staged multi package install of APEX and APK packages is not supported");
-            }
             List<VerificationParams> verifyingChildSessions =
                     new ArrayList<>(childSessions.size());
             boolean success = true;
@@ -4172,6 +4165,12 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
                     + childSession.sessionId + " and session " + sessionId
                     + " have inconsistent rollback settings");
         }
+        boolean hasAPK = containsApkSession() || !childSession.isApexSession();
+        boolean hasAPEX = sessionContains(s -> s.isApexSession()) || childSession.isApexSession();
+        if (!params.isStaged && hasAPK && hasAPEX) {
+            throw new IllegalStateException("Mix of APK and APEX is not supported for "
+                    + "non-staged multi-package session");
+        }
 
         try {
             acquireTransactionLock();
diff --git a/services/core/java/com/android/server/pm/VerificationParams.java b/services/core/java/com/android/server/pm/VerificationParams.java
index 3d2ffe1b868e9..c5569cd53da10 100644
--- a/services/core/java/com/android/server/pm/VerificationParams.java
+++ b/services/core/java/com/android/server/pm/VerificationParams.java
@@ -65,7 +65,7 @@ import android.os.UserHandle;
 import android.os.UserManager;
 import android.provider.DeviceConfig;
 import android.provider.Settings;
-import android.util.ArrayMap;
+import android.util.ArraySet;
 import android.util.Pair;
 import android.util.Slog;
 
@@ -78,7 +78,7 @@ import java.security.cert.CertificateException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
+import java.util.Set;
 
 final class VerificationParams extends HandlerParams {
     /**
@@ -688,7 +688,7 @@ final class VerificationParams extends HandlerParams {
 
     private void sendVerificationCompleteNotification() {
         if (mParentVerificationParams != null) {
-            mParentVerificationParams.trySendVerificationCompleteNotification(this, mRet);
+            mParentVerificationParams.trySendVerificationCompleteNotification(this);
         } else {
             try {
                 mObserver.onPackageInstalled(null, mRet, mErrorMessage,
@@ -717,7 +717,7 @@ final class VerificationParams extends HandlerParams {
     static final class MultiPackageVerificationParams extends HandlerParams {
         private final IPackageInstallObserver2 mObserver;
         private final List<VerificationParams> mChildParams;
-        private final Map<VerificationParams, Integer> mVerificationState;
+        private final Set<VerificationParams> mVerificationState;
 
         MultiPackageVerificationParams(VerificationParams parent, List<VerificationParams> children,
                 PackageManagerService pm) throws PackageManagerException {
@@ -731,7 +731,7 @@ final class VerificationParams extends HandlerParams {
                 final VerificationParams childParams = children.get(i);
                 childParams.mParentVerificationParams = this;
             }
-            mVerificationState = new ArrayMap<>(mChildParams.size());
+            mVerificationState = new ArraySet<>(mChildParams.size());
             mObserver = parent.mObserver;
         }
 
@@ -749,18 +749,16 @@ final class VerificationParams extends HandlerParams {
             }
         }
 
-        void trySendVerificationCompleteNotification(VerificationParams child, int currentStatus) {
-            mVerificationState.put(child, currentStatus);
+        void trySendVerificationCompleteNotification(VerificationParams child) {
+            mVerificationState.add(child);
             if (mVerificationState.size() != mChildParams.size()) {
                 return;
             }
             int completeStatus = PackageManager.INSTALL_SUCCEEDED;
             String errorMsg = null;
-            for (VerificationParams params : mVerificationState.keySet()) {
+            for (VerificationParams params : mVerificationState) {
                 int status = params.mRet;
-                if (status == PackageManager.INSTALL_UNKNOWN) {
-                    return;
-                } else if (status != PackageManager.INSTALL_SUCCEEDED) {
+                if (status != PackageManager.INSTALL_SUCCEEDED) {
                     completeStatus = status;
                     errorMsg = params.mErrorMessage;
                     break;