diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 1648e7f66c721..efde6f1ea2549 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -7273,8 +7273,10 @@ package android.os {
     method public boolean bind(android.os.UpdateEngineCallback);
     method public void cancel();
     method @WorkerThread public int cleanupAppliedPayload();
+    method public void resetShouldSwitchSlotOnReboot();
     method public void resetStatus();
     method public void resume();
+    method public void setShouldSwitchSlotOnReboot(@NonNull String);
     method public void suspend();
     method public boolean unbind();
     method public boolean verifyPayloadMetadata(String);
diff --git a/core/java/android/os/UpdateEngine.java b/core/java/android/os/UpdateEngine.java
index 5a48242f26eae..3e01c53f04693 100644
--- a/core/java/android/os/UpdateEngine.java
+++ b/core/java/android/os/UpdateEngine.java
@@ -409,10 +409,11 @@ public class UpdateEngine {
 
     /**
      * Resets the bootable flag on the non-current partition and all internal
-     * update_engine state. This can be used after an unwanted payload has been
-     * successfully applied and the device has not yet been rebooted to signal
-     * that we no longer want to boot into that updated system. After this call
-     * completes, update_engine will no longer report
+     * update_engine state. Note this call will clear the entire update
+     * progress. So a subsequent {@link #applyPayload} will apply the update
+     * from scratch.
+     *
+     * <p>After this call completes, update_engine will no longer report
      * {@code UPDATED_NEED_REBOOT}, so your callback can remove any outstanding
      * notification that rebooting into the new system is possible.
      */
@@ -424,6 +425,39 @@ public class UpdateEngine {
         }
     }
 
+    /**
+     * Sets the A/B slot switch for the next boot after applying an ota update. If
+     * {@link #applyPayload} hasn't switched the slot, the updater APP can call
+     * this API to switch the slot and apply the update on next boot.
+     *
+     * @param payloadMetadataFilename the location of the metadata without the
+     * {@code file://} prefix.
+     */
+    public void setShouldSwitchSlotOnReboot(@NonNull String payloadMetadataFilename) {
+        try {
+            mUpdateEngine.setShouldSwitchSlotOnReboot(payloadMetadataFilename);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+   /**
+    * Resets the boot slot to the source/current slot, without cancelling the
+    * update progress. This can be called after the update is installed, and to
+    * prevent the device from accidentally taking the update when it reboots.
+    *
+    * This is useful when users don't want to take the update immediately; or
+    * the updater determines some condition hasn't met, e.g. insufficient space
+    * for boot.
+    */
+    public void resetShouldSwitchSlotOnReboot() {
+        try {
+            mUpdateEngine.resetShouldSwitchSlotOnReboot();
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
     /**
      * Unbinds the last bound callback function.
      */