diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index f9981d0d9c968..10b9bb4cdf24c 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -205,14 +205,16 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                             PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_INSTALL_FLAGS, 0);
                     int user = intent.getIntExtra(
                             PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_USER, 0);
+                    int sessionId = intent.getIntExtra(
+                            PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_SESSION_ID, -1);
 
                     File newPackageCodePath = new File(intent.getData().getPath());
 
                     queueSleepIfNeeded();
 
                     getHandler().post(() -> {
-                        boolean success =
-                                enableRollback(installFlags, newPackageCodePath, user, token);
+                        boolean success = enableRollback(
+                                sessionId, installFlags, newPackageCodePath, user, token);
                         int ret = PackageManagerInternal.ENABLE_ROLLBACK_SUCCEEDED;
                         if (!success) {
                             ret = PackageManagerInternal.ENABLE_ROLLBACK_FAILED;
@@ -240,17 +242,16 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                 if (Intent.ACTION_CANCEL_ENABLE_ROLLBACK.equals(intent.getAction())) {
                     int token = intent.getIntExtra(
                             PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_TOKEN, -1);
+                    int sessionId = intent.getIntExtra(
+                            PackageManagerInternal.EXTRA_ENABLE_ROLLBACK_SESSION_ID, -1);
                     if (LOCAL_LOGV) {
                         Slog.v(TAG, "broadcast=ACTION_CANCEL_ENABLE_ROLLBACK token=" + token);
                     }
                     synchronized (mLock) {
-                        for (int i = 0; i < mRollbacks.size(); ++i) {
-                            Rollback rollback = mRollbacks.get(i);
-                            if (rollback.hasToken(token) && rollback.isEnabling()) {
-                                mRollbacks.remove(i);
-                                rollback.delete(mAppDataRollbackHelper);
-                                break;
-                            }
+                        Rollback rollback = getRollbackForSessionLocked(sessionId);
+                        if (rollback != null && rollback.isEnabling()) {
+                            mRollbacks.remove(rollback);
+                            rollback.delete(mAppDataRollbackHelper);
                         }
                     }
                 }
@@ -723,51 +724,24 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
      * @return true if enabling the rollback succeeds, false otherwise.
      */
     @WorkerThread
-    private boolean enableRollback(
+    private boolean enableRollback(int sessionId,
             int installFlags, File newPackageCodePath, @UserIdInt int user, int token) {
         if (LOCAL_LOGV) {
             Slog.v(TAG, "enableRollback user=" + user + " token=" + token
                     + " path=" + newPackageCodePath.getAbsolutePath());
         }
 
-        // Find the session id associated with this install.
-        // TODO: It would be nice if package manager or package installer told
-        // us the session directly, rather than have to search for it
-        // ourselves.
-
-        // getAllSessions only returns sessions for the associated user.
-        // Create a context with the right user so we can find the matching
-        // session.
-        final Context context = getContextAsUser(UserHandle.of(user));
-        if (context == null) {
-            Slog.e(TAG, "Unable to create context for install session user.");
+        PackageInstaller installer = mContext.getPackageManager().getPackageInstaller();
+        PackageInstaller.SessionInfo packageSession = installer.getSessionInfo(sessionId);
+        if (packageSession == null) {
+            Slog.e(TAG, "Unable to find session for enabled rollback.");
             return false;
         }
 
-        PackageInstaller.SessionInfo parentSession = null;
-        PackageInstaller.SessionInfo packageSession = null;
-        PackageInstaller installer = context.getPackageManager().getPackageInstaller();
-        for (PackageInstaller.SessionInfo info : installer.getAllSessions()) {
-            if (info.isMultiPackage()) {
-                for (int childId : info.getChildSessionIds()) {
-                    PackageInstaller.SessionInfo child = installer.getSessionInfo(childId);
-                    if (sessionMatchesForEnableRollback(child, installFlags, newPackageCodePath)) {
-                        // TODO: Check we only have one matching session?
-                        parentSession = info;
-                        packageSession = child;
-                        break;
-                    }
-                }
-            } else if (sessionMatchesForEnableRollback(info, installFlags, newPackageCodePath)) {
-                // TODO: Check we only have one matching session?
-                parentSession = info;
-                packageSession = info;
-                break;
-            }
-        }
-
-        if (parentSession == null || packageSession == null) {
-            Slog.e(TAG, "Unable to find session for enabled rollback.");
+        PackageInstaller.SessionInfo parentSession = packageSession.hasParentSessionId()
+                ? installer.getSessionInfo(packageSession.getParentSessionId()) : packageSession;
+        if (parentSession == null) {
+            Slog.e(TAG, "Unable to find parent session for enabled rollback.");
             return false;
         }