diff --git a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
index be842b9..9080d46 100644
--- a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
+++ b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
@@ -86,7 +86,6 @@
 import android.os.ServiceSpecificException;
 import android.os.ShellCommand;
 import android.os.SystemClock;
-import android.os.SystemProperties;
 import android.os.Trace;
 import android.os.UserHandle;
 import android.os.UserManager;
@@ -3053,6 +3052,13 @@
         // Set package source to other by default
         sessionParams.setPackageSource(PackageInstaller.PACKAGE_SOURCE_OTHER);
 
+        // Encodes one of the states:
+        //  1. Install request explicitly specified --staged, then value will be true.
+        //  2. Install request explicitly specified --non-staged, then value will be false.
+        //  3. Install request did not specify either --staged or --non-staged, then for APEX
+        //      installs the value will be true, and for apk installs it will be false.
+        Boolean staged = null;
+
         String opt;
         boolean replaceExisting = true;
         boolean forceNonStaged = false;
@@ -3151,7 +3157,6 @@
                     break;
                 case "--apex":
                     sessionParams.setInstallAsApex();
-                    sessionParams.setStaged();
                     break;
                 case "--force-non-staged":
                     forceNonStaged = true;
@@ -3160,7 +3165,10 @@
                     sessionParams.setMultiPackage();
                     break;
                 case "--staged":
-                    sessionParams.setStaged();
+                    staged = true;
+                    break;
+                case "--non-staged":
+                    staged = false;
                     break;
                 case "--force-queryable":
                     sessionParams.setForceQueryable();
@@ -3192,11 +3200,16 @@
                     throw new IllegalArgumentException("Unknown option " + opt);
             }
         }
+        if (staged == null) {
+            staged = (sessionParams.installFlags & PackageManager.INSTALL_APEX) != 0;
+        }
         if (replaceExisting) {
             sessionParams.installFlags |= PackageManager.INSTALL_REPLACE_EXISTING;
         }
         if (forceNonStaged) {
             sessionParams.isStaged = false;
+        } else if (staged) {
+            sessionParams.setStaged();
         }
         return params;
     }
@@ -3978,7 +3991,8 @@
         pw.println("       [--preload] [--instant] [--full] [--dont-kill]");
         pw.println("       [--enable-rollback]");
         pw.println("       [--force-uuid internal|UUID] [--pkg PACKAGE] [-S BYTES]");
-        pw.println("       [--apex] [--force-non-staged] [--staged-ready-timeout TIMEOUT]");
+        pw.println("       [--apex] [--non-staged] [--force-non-staged]");
+        pw.println("       [--staged-ready-timeout TIMEOUT]");
         pw.println("       [PATH [SPLIT...]|-]");
         pw.println("    Install an application.  Must provide the apk data to install, either as");
         pw.println("    file path(s) or '-' to read from stdin.  Options are:");
@@ -4006,6 +4020,9 @@
         pw.println("          3=device setup, 4=user request");
         pw.println("      --force-uuid: force install on to disk volume with given UUID");
         pw.println("      --apex: install an .apex file, not an .apk");
+        pw.println("      --non-staged: explicitly set this installation to be non-staged.");
+        pw.println("          This flag is only useful for APEX installs that are implicitly");
+        pw.println("          assumed to be staged.");
         pw.println("      --force-non-staged: force the installation to run under a non-staged");
         pw.println("          session, which may complete without requiring a reboot");
         pw.println("      --staged-ready-timeout: By default, staged sessions wait "