diff --git a/services/core/java/com/android/server/pm/StagingManager.java b/services/core/java/com/android/server/pm/StagingManager.java
index 14da9aa144706..616e5d13f9905 100644
--- a/services/core/java/com/android/server/pm/StagingManager.java
+++ b/services/core/java/com/android/server/pm/StagingManager.java
@@ -738,7 +738,6 @@ public class StagingManager {
         PackageInstaller.SessionParams params = originalSession.params.copy();
         params.isStaged = false;
         params.installFlags |= PackageManager.INSTALL_STAGED;
-        // TODO(b/129744602): use the userid from the original session.
         if (preReboot) {
             params.installFlags &= ~PackageManager.INSTALL_ENABLE_ROLLBACK;
             params.installFlags |= PackageManager.INSTALL_DRY_RUN;
@@ -748,7 +747,7 @@ public class StagingManager {
         try {
             int apkSessionId = mPi.createSession(
                     params, originalSession.getInstallerPackageName(),
-                    originalSession.getInstallerAttributionTag(), 0 /* UserHandle.SYSTEM */);
+                    originalSession.getInstallerAttributionTag(), originalSession.userId);
             PackageInstallerSession apkSession = mPi.getSession(apkSessionId);
             apkSession.open();
             for (int i = 0, size = apkFilePaths.size(); i < size; i++) {
@@ -806,10 +805,9 @@ public class StagingManager {
             if (preReboot) {
                 params.installFlags &= ~PackageManager.INSTALL_ENABLE_ROLLBACK;
             }
-            // TODO(b/129744602): use the userid from the original session.
             final int apkParentSessionId = mPi.createSession(
                     params, session.getInstallerPackageName(), session.getInstallerAttributionTag(),
-                    0 /* UserHandle.SYSTEM */);
+                    session.userId);
             final PackageInstallerSession apkParentSession = mPi.getSession(apkParentSessionId);
             try {
                 apkParentSession.open();
diff --git a/tests/RollbackTest/Android.bp b/tests/RollbackTest/Android.bp
index a23df920b3968..4f5a30502c913 100644
--- a/tests/RollbackTest/Android.bp
+++ b/tests/RollbackTest/Android.bp
@@ -53,6 +53,9 @@ java_test_host {
     name: "MultiUserRollbackTest",
     srcs: ["MultiUserRollbackTest/src/**/*.java"],
     libs: ["tradefed"],
+    static_libs: [
+        "frameworks-base-hostutils",
+    ],
     test_suites: ["general-tests"],
     test_config: "MultiUserRollbackTest.xml",
 }
diff --git a/tests/RollbackTest/MultiUserRollbackTest/src/com/android/tests/rollback/host/MultiUserRollbackTest.java b/tests/RollbackTest/MultiUserRollbackTest/src/com/android/tests/rollback/host/MultiUserRollbackTest.java
index 42b886f0774f7..f160847448537 100644
--- a/tests/RollbackTest/MultiUserRollbackTest/src/com/android/tests/rollback/host/MultiUserRollbackTest.java
+++ b/tests/RollbackTest/MultiUserRollbackTest/src/com/android/tests/rollback/host/MultiUserRollbackTest.java
@@ -24,6 +24,7 @@ import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -40,6 +41,9 @@ public class MultiUserRollbackTest extends BaseHostJUnit4Test {
     private static final long SWITCH_USER_COMPLETED_NUMBER_OF_POLLS = 60;
     private static final long SWITCH_USER_COMPLETED_POLL_INTERVAL_IN_MILLIS = 1000;
 
+    @Rule
+    public AbandonSessionsRule mHostTestRule = new AbandonSessionsRule(this);
+
     @After
     public void tearDown() throws Exception {
         removeSecondaryUserIfNecessary();
@@ -59,6 +63,30 @@ public class MultiUserRollbackTest extends BaseHostJUnit4Test {
         runPhaseForUsers("testBasic", mSecondaryUserId);
     }
 
+    /**
+     * Tests staged install/rollback works correctly on the 2nd user.
+     */
+    @Test
+    public void testStagedRollback() throws Exception {
+        runPhaseForUsers("testStagedRollback_Phase1", mSecondaryUserId);
+        getDevice().reboot();
+
+        // Need to unlock the user for device tests to run successfully
+        getDevice().startUser(mSecondaryUserId);
+        awaitUserUnlocked(mSecondaryUserId);
+        runPhaseForUsers("testStagedRollback_Phase2", mSecondaryUserId);
+        getDevice().reboot();
+
+        getDevice().startUser(mSecondaryUserId);
+        awaitUserUnlocked(mSecondaryUserId);
+        runPhaseForUsers("testStagedRollback_Phase3", mSecondaryUserId);
+        getDevice().reboot();
+
+        getDevice().startUser(mSecondaryUserId);
+        awaitUserUnlocked(mSecondaryUserId);
+        runPhaseForUsers("testStagedRollback_Phase4", mSecondaryUserId);
+    }
+
     @Test
     public void testMultipleUsers() throws Exception {
         runPhaseForUsers("testMultipleUsersInstallV1", mOriginalUserId, mSecondaryUserId);
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java
index 8641f4d4013ab..5d133a4de13dc 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/MultiUserRollbackTest.java
@@ -115,4 +115,32 @@ public class MultiUserRollbackTest {
         assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
         InstallUtils.processUserData(TestApp.A);
     }
+
+    @Test
+    public void testStagedRollback_Phase1() throws Exception {
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(-1);
+        Install.single(TestApp.A1).setStaged().commit();
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(-1);
+    }
+
+    @Test
+    public void testStagedRollback_Phase2() throws Exception {
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
+        Install.single(TestApp.A2).setStaged().setEnableRollback().commit();
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
+    }
+
+    @Test
+    public void testStagedRollback_Phase3() throws Exception {
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(2);
+        RollbackInfo rollback = RollbackUtils.waitForAvailableRollback(TestApp.A);
+        assertThat(rollback).packagesContainsExactly(Rollback.from(TestApp.A2).to(TestApp.A1));
+        RollbackUtils.rollback(rollback.getRollbackId());
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(2);
+    }
+
+    @Test
+    public void testStagedRollback_Phase4() {
+        assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
+    }
 }