diff --git a/services/core/java/com/android/server/rollback/RollbackData.java b/services/core/java/com/android/server/rollback/RollbackData.java
index f456665..fcd5297 100644
--- a/services/core/java/com/android/server/rollback/RollbackData.java
+++ b/services/core/java/com/android/server/rollback/RollbackData.java
@@ -64,6 +64,11 @@
     public boolean isAvailable;
 
     /**
+     * The id of the post-reboot apk session for a staged install, if any.
+     */
+    public int apkSessionId = -1;
+
+    /**
      * Whether this Rollback is currently in progress. This field is true from the point
      * we commit a {@code PackageInstaller} session containing these packages to the point the
      * {@code PackageInstaller} calls into the {@code onFinished} callback.
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 7445c00..dfb95b8 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -930,7 +930,28 @@
 
     @Override
     public void notifyStagedApkSession(int originalSessionId, int apkSessionId) {
-        // TODO: Handle this callback.
+        getHandler().post(() -> {
+            RollbackData rd = null;
+            synchronized (mLock) {
+                ensureRollbackDataLoadedLocked();
+                for (int i = 0; i < mAvailableRollbacks.size(); ++i) {
+                    RollbackData data = mAvailableRollbacks.get(i);
+                    if (data.stagedSessionId == originalSessionId) {
+                        data.apkSessionId = apkSessionId;
+                        rd = data;
+                        break;
+                    }
+                }
+            }
+
+            if (rd != null) {
+                try {
+                    mRollbackStore.saveAvailableRollback(rd);
+                } catch (IOException ioe) {
+                    Log.e(TAG, "Unable to save rollback info for : " + rd.rollbackId, ioe);
+                }
+            }
+        });
     }
 
     /**
diff --git a/services/core/java/com/android/server/rollback/RollbackStore.java b/services/core/java/com/android/server/rollback/RollbackStore.java
index 1e45ba2..d17ebae 100644
--- a/services/core/java/com/android/server/rollback/RollbackStore.java
+++ b/services/core/java/com/android/server/rollback/RollbackStore.java
@@ -248,6 +248,7 @@
             dataJson.put("timestamp", data.timestamp.toString());
             dataJson.put("stagedSessionId", data.stagedSessionId);
             dataJson.put("isAvailable", data.isAvailable);
+            dataJson.put("apkSessionId", data.apkSessionId);
 
             PrintWriter pw = new PrintWriter(new File(data.backupDir, "rollback.json"));
             pw.println(dataJson.toString());
@@ -313,6 +314,7 @@
                     stagedSessionId, isAvailable);
             data.packages.addAll(packageRollbackInfosFromJson(dataJson.getJSONArray("packages")));
             data.timestamp = Instant.parse(dataJson.getString("timestamp"));
+            data.apkSessionId = dataJson.getInt("apkSessionId");
             return data;
         } catch (JSONException | DateTimeParseException e) {
             throw new IOException(e);