diff --git a/aosp/update_engine_client_android.cc b/aosp/update_engine_client_android.cc
index 1a68cf4..8628463 100644
--- a/aosp/update_engine_client_android.cc
+++ b/aosp/update_engine_client_android.cc
@@ -41,6 +41,7 @@
 #include "update_engine/common/error_code.h"
 #include "update_engine/common/error_code_utils.h"
 #include "update_engine/update_status_utils.h"
+#include "utils/String8.h"
 
 using android::binder::Status;
 
@@ -71,6 +72,8 @@
 
   // Called whenever the UpdateEngine daemon dies.
   void UpdateEngineServiceDied();
+  // Register callback to watch for death notification from update_engine.
+  void RegisterDeathNotification();
 
   static std::vector<android::String16> ParseHeaders(const std::string& arg);
 
@@ -106,6 +109,18 @@
   return Status::ok();
 }
 
+constexpr auto&& UNSPECIFIED_FLAG = "unspecified";
+
+void UpdateEngineClientAndroid::RegisterDeathNotification() {
+  // When following updates status changes, exit if the update_engine daemon
+  // dies.
+  android::BinderWrapper::Create();
+  android::BinderWrapper::Get()->RegisterForDeathNotifications(
+      android::os::IUpdateEngine::asBinder(service_),
+      base::Bind(&UpdateEngineClientAndroid::UpdateEngineServiceDied,
+                 base::Unretained(this)));
+}
+
 int UpdateEngineClientAndroid::OnInit() {
   int ret = Daemon::OnInit();
   if (ret != EX_OK)
@@ -137,6 +152,11 @@
                 "The path to the update payload metadata. "
                 "Used when --verify or --allocate is passed.");
 
+  DEFINE_string(switch_slot,
+                UNSPECIFIED_FLAG,
+                "Perform just the slow switching part of OTA. "
+                "Used to revert a slot switch or re-do slot switch. Valid "
+                "values are 'true' and 'false'");
   DEFINE_bool(suspend, false, "Suspend an ongoing update and exit.");
   DEFINE_bool(resume, false, "Resume a suspended update.");
   DEFINE_bool(cancel, false, "Cancel the ongoing update and exit.");
@@ -198,6 +218,23 @@
     return ExitWhenIdle(service_->resetStatus());
   }
 
+  if (FLAGS_switch_slot != UNSPECIFIED_FLAG) {
+    if (FLAGS_switch_slot != "true" && FLAGS_switch_slot != "false") {
+      LOG(ERROR) << "--switch_slot should be either true or false, got "
+                 << FLAGS_switch_slot;
+      return 1;
+    }
+    const bool should_switch = FLAGS_switch_slot == "true";
+    ::android::binder::Status status;
+    if (should_switch) {
+      status = service_->setShouldSwitchSlotOnReboot(
+          android::String16(FLAGS_metadata.c_str(), FLAGS_metadata.size()));
+    } else {
+      status = service_->resetShouldSwitchSlotOnReboot();
+    }
+    return ExitWhenIdle(status);
+  }
+
   if (FLAGS_verify) {
     bool applicable = false;
     Status status = service_->verifyPayloadApplicable(
@@ -262,14 +299,7 @@
   if (!keep_running)
     return ExitWhenIdle(EX_OK);
 
-  // When following updates status changes, exit if the update_engine daemon
-  // dies.
-  android::BinderWrapper::Create();
-  android::BinderWrapper::Get()->RegisterForDeathNotifications(
-      android::os::IUpdateEngine::asBinder(service_),
-      base::Bind(&UpdateEngineClientAndroid::UpdateEngineServiceDied,
-                 base::Unretained(this)));
-
+  RegisterDeathNotification();
   return EX_OK;
 }