diff --git a/services/core/java/com/android/server/pm/PackageInstallerSession.java b/services/core/java/com/android/server/pm/PackageInstallerSession.java
index e0f1b0b44cf8a..4d0d5532399ec 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerSession.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerSession.java
@@ -2261,10 +2261,6 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
         }
 
         if (params.isStaged) {
-            // TODO(b/136257624): CTS test fails if we don't send session finished broadcast, even
-            //  though ideally, we just need to send session committed broadcast.
-            sendUpdateToRemoteStatusReceiver(INSTALL_SUCCEEDED, "Session staged", null);
-
             mStagedSession.verifySession();
         } else {
             verify();
@@ -2511,6 +2507,7 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
                 mStagedSession.notifyEndPreRebootVerification();
                 if (error == SessionInfo.SESSION_NO_ERROR) {
                     mStagingManager.commitSession(mStagedSession);
+                    sendUpdateToRemoteStatusReceiver(INSTALL_SUCCEEDED, "Session staged", null);
                 } else {
                     dispatchSessionFinished(INSTALL_FAILED_VERIFICATION_FAILURE, msg, null);
                     maybeFinishChildSessions(INSTALL_FAILED_VERIFICATION_FAILURE, msg);
diff --git a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
index fd2256f149639..99f70b206b65c 100644
--- a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
+++ b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
@@ -3563,18 +3563,27 @@ class PackageManagerShellCommand extends ShellCommand {
             }
             final LocalIntentReceiver receiver = new LocalIntentReceiver();
             session.commit(receiver.getIntentSender());
-            final Intent result = receiver.getResult();
-            final int status = result.getIntExtra(PackageInstaller.EXTRA_STATUS,
-                    PackageInstaller.STATUS_FAILURE);
-            if (status == PackageInstaller.STATUS_SUCCESS) {
+            if (!session.isStaged()) {
+                final Intent result = receiver.getResult();
+                final int status = result.getIntExtra(PackageInstaller.EXTRA_STATUS,
+                        PackageInstaller.STATUS_FAILURE);
+                if (status == PackageInstaller.STATUS_SUCCESS) {
+                    if (logSuccess) {
+                        pw.println("Success");
+                    }
+                } else {
+                    pw.println("Failure ["
+                            + result.getStringExtra(PackageInstaller.EXTRA_STATUS_MESSAGE) + "]");
+                }
+                return status;
+            } else {
+                // Return immediately without retrieving the result. The caller will decide
+                // whether to wait for the session to become ready.
                 if (logSuccess) {
                     pw.println("Success");
                 }
-            } else {
-                pw.println("Failure ["
-                        + result.getStringExtra(PackageInstaller.EXTRA_STATUS_MESSAGE) + "]");
+                return PackageInstaller.STATUS_SUCCESS;
             }
-            return status;
         } finally {
             IoUtils.closeQuietly(session);
         }