diff --git a/services/core/java/com/android/server/PackageWatchdog.java b/services/core/java/com/android/server/PackageWatchdog.java
index 4507193..ec40971 100644
--- a/services/core/java/com/android/server/PackageWatchdog.java
+++ b/services/core/java/com/android/server/PackageWatchdog.java
@@ -247,9 +247,8 @@
                 }
 
                 for (int pIndex = 0; pIndex < packages.size(); pIndex++) {
-                    String packageToReport = packages.get(pIndex).getPackageName();
-                    long packageVersionCode = packages.get(pIndex).getVersionCode();
-                    // Observer that will receive failure for packageToReport
+                    VersionedPackage versionedPackage = packages.get(pIndex);
+                    // Observer that will receive failure for versionedPackage
                     PackageHealthObserver currentObserverToNotify = null;
                     int currentObserverImpact = Integer.MAX_VALUE;
 
@@ -258,9 +257,8 @@
                         ObserverInternal observer = mAllObservers.valueAt(oIndex);
                         PackageHealthObserver registeredObserver = observer.mRegisteredObserver;
                         if (registeredObserver != null
-                                && observer.onPackageFailure(packageToReport)) {
-                            int impact = registeredObserver.onHealthCheckFailed(packageToReport,
-                                    packageVersionCode);
+                                && observer.onPackageFailure(versionedPackage.getPackageName())) {
+                            int impact = registeredObserver.onHealthCheckFailed(versionedPackage);
                             if (impact != PackageHealthObserverImpact.USER_IMPACT_NONE
                                     && impact < currentObserverImpact) {
                                 currentObserverToNotify = registeredObserver;
@@ -271,7 +269,7 @@
 
                     // Execute action with least user impact
                     if (currentObserverToNotify != null) {
-                        currentObserverToNotify.execute(packageToReport, packageVersionCode);
+                        currentObserverToNotify.execute(versionedPackage);
                     }
                 }
             }
@@ -310,19 +308,19 @@
     /** Register instances of this interface to receive notifications on package failure. */
     public interface PackageHealthObserver {
         /**
-         * Called when health check fails for the {@code packageName}.
+         * Called when health check fails for the {@code versionedPackage}.
          *
          * @return any one of {@link PackageHealthObserverImpact} to express the impact
          * to the user on {@link #execute}
          */
-        @PackageHealthObserverImpact int onHealthCheckFailed(String packageName, long versionCdoe);
+        @PackageHealthObserverImpact int onHealthCheckFailed(VersionedPackage versionedPackage);
 
         /**
          * Executes mitigation for {@link #onHealthCheckFailed}.
          *
          * @return {@code true} if action was executed successfully, {@code false} otherwise
          */
-        boolean execute(String packageName, long versionCode);
+        boolean execute(VersionedPackage versionedPackage);
 
         // TODO(zezeozue): Ensure uniqueness?
         /**
diff --git a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
index d00817f..8e04160 100644
--- a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
+++ b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
@@ -57,7 +57,7 @@
     }
 
     @Override
-    public int onHealthCheckFailed(String packageName, long versionCode) {
+    public int onHealthCheckFailed(VersionedPackage failedPackage) {
         VersionedPackage moduleMetadataPackage = getModuleMetadataPackage();
         if (moduleMetadataPackage == null) {
             // Ignore failure, no mainline update available
@@ -66,7 +66,7 @@
 
         RollbackInfo rollback =
                 getAvailableMainlineRollback(mContext.getSystemService(RollbackManager.class),
-                    packageName, versionCode, moduleMetadataPackage);
+                        failedPackage, moduleMetadataPackage);
         if (rollback == null) {
             // Don't handle the notification, no rollbacks available for the package
             return PackageHealthObserverImpact.USER_IMPACT_NONE;
@@ -76,7 +76,7 @@
     }
 
     @Override
-    public boolean execute(String packageName, long versionCode) {
+    public boolean execute(VersionedPackage failedPackage) {
         VersionedPackage moduleMetadataPackage = getModuleMetadataPackage();
         if (moduleMetadataPackage == null) {
             // Ignore failure, no mainline update available
@@ -85,10 +85,11 @@
 
         RollbackManager rollbackManager = mContext.getSystemService(RollbackManager.class);
         RollbackInfo rollback = getAvailableMainlineRollback(rollbackManager,
-                packageName, versionCode, moduleMetadataPackage);
+                failedPackage, moduleMetadataPackage);
         if (rollback == null) {
-            Slog.w(TAG, "Expected rollback but no rollback found for package: [ "
-                    + packageName + "] with versionCode: [" + versionCode + "]");
+            Slog.w(TAG, "Expected rollback but no mainline rollback found for package: [ "
+                    + failedPackage.getPackageName() + "] with versionCode: ["
+                    + failedPackage.getVersionCode() + "]");
             return false;
         }
 
@@ -114,7 +115,7 @@
 
         mHandler.post(() ->
                 rollbackManager.commitRollback(rollback.getRollbackId(),
-                    Collections.singletonList(new VersionedPackage(packageName, versionCode)),
+                    Collections.singletonList(moduleMetadataPackage),
                     rollbackReceiver.getIntentSender()));
         // Assume rollback executed successfully
         return true;
@@ -134,7 +135,7 @@
     }
 
     private RollbackInfo getAvailableMainlineRollback(RollbackManager rollbackManager,
-            String packageName, long versionCode, VersionedPackage moduleMetadataPackage) {
+            VersionedPackage failedPackage, VersionedPackage moduleMetadataPackage) {
         for (RollbackInfo rollback : rollbackManager.getAvailableRollbacks()) {
             // We only rollback mainline packages, so check if rollback contains the
             // module metadata provider, if it does, the rollback is a mainline rollback
@@ -142,10 +143,11 @@
             boolean hasFailedPackage = false;
             for (PackageRollbackInfo packageRollback : rollback.getPackages()) {
                 hasModuleMetadataPackage |= packageRollback.getPackageName().equals(
-                                moduleMetadataPackage.getPackageName());
-                hasFailedPackage |= packageRollback.getPackageName().equals(packageName)
+                        moduleMetadataPackage.getPackageName());
+                hasFailedPackage |= packageRollback.getPackageName().equals(
+                        failedPackage.getPackageName())
                         && packageRollback.getVersionRolledBackFrom().getVersionCode()
-                        == versionCode;
+                        == failedPackage.getVersionCode();
             }
             if (hasModuleMetadataPackage && hasFailedPackage) {
                 return rollback;
diff --git a/tests/PackageWatchdog/src/com/android/server/PackageWatchdogTest.java b/tests/PackageWatchdog/src/com/android/server/PackageWatchdogTest.java
index c1c598d..77cd62e 100644
--- a/tests/PackageWatchdog/src/com/android/server/PackageWatchdogTest.java
+++ b/tests/PackageWatchdog/src/com/android/server/PackageWatchdogTest.java
@@ -242,8 +242,9 @@
         PackageWatchdog watchdog = createWatchdog();
         long differentVersionCode = 2L;
         TestObserver observer = new TestObserver(OBSERVER_NAME_1) {
-                public int onHealthCheckFailed(String packageName, long versionCode) {
-                    if (versionCode == VERSION_CODE) {
+                @Override
+                public int onHealthCheckFailed(VersionedPackage versionedPackage) {
+                    if (versionedPackage.getVersionCode() == VERSION_CODE) {
                         // Only rollback for specific versionCode
                         return PackageHealthObserverImpact.USER_IMPACT_MEDIUM;
                     }
@@ -457,12 +458,12 @@
             mImpact = impact;
         }
 
-        public int onHealthCheckFailed(String packageName, long versionCode) {
+        public int onHealthCheckFailed(VersionedPackage versionedPackage) {
             return mImpact;
         }
 
-        public boolean execute(String packageName, long versionCode) {
-            mFailedPackages.add(packageName);
+        public boolean execute(VersionedPackage versionedPackage) {
+            mFailedPackages.add(versionedPackage.getPackageName());
             return true;
         }