18 files changed, 507 insertions, 206 deletions

diff --git a/common.h b/common.h
index 3dc36a960..c24431bd1 100644
--- a/common.h
+++ b/common.h
@@ -32,6 +32,7 @@ struct selabel_handle;
 extern struct selabel_handle* sehandle;
 extern RecoveryUI* ui;
 extern bool modified_flash;
+extern bool has_cache;
 
 // The current stage, e.g. "1/2".
 extern std::string stage;
diff --git a/recovery.cpp b/recovery.cpp
index f03cec3fe..ac3e7c633 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "private/recovery.h"
+#include "recovery.h"
 
 #include <ctype.h>
 #include <dirent.h>
@@ -32,7 +32,6 @@
 #include <sys/stat.h>
 #include <sys/types.h>
 #include <sys/wait.h>
-#include <time.h>
 #include <unistd.h>
 
 #include <algorithm>
@@ -49,12 +48,8 @@
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 #include <bootloader_message/bootloader_message.h>
-#include <cutils/android_reboot.h>
 #include <cutils/properties.h> /* for property_list */
 #include <health2/Health.h>
-#include <selinux/android.h>
-#include <selinux/label.h>
-#include <selinux/selinux.h>
 #include <ziparchive/zip_archive.h>
 
 #include "adb_install.h"
@@ -70,7 +65,6 @@
 #include "otautil/sysutil.h"
 #include "roots.h"
 #include "screen_ui.h"
-#include "stub_ui.h"
 #include "ui.h"
 
 static constexpr const char* CACHE_LOG_DIR = "/cache/recovery";
@@ -88,13 +82,9 @@ static constexpr const char* SDCARD_ROOT = "/sdcard";
 // into target_files.zip. Assert the version defined in code and in Android.mk are consistent.
 static_assert(kRecoveryApiVersion == RECOVERY_API_VERSION, "Mismatching recovery API versions.");
 
-static bool has_cache = false;
-
-RecoveryUI* ui = nullptr;
 bool modified_flash = false;
 std::string stage;
 const char* reason = nullptr;
-struct selabel_handle* sehandle;
 
 /*
  * The recovery tool communicates with the main system through /cache files.
@@ -146,77 +136,6 @@ bool is_ro_debuggable() {
     return android::base::GetBoolProperty("ro.debuggable", false);
 }
 
-// command line args come from, in decreasing precedence:
-//   - the actual command line
-//   - the bootloader control block (one per line, after "recovery")
-//   - the contents of COMMAND_FILE (one per line)
-static std::vector<std::string> get_args(const int argc, char** const argv) {
-  CHECK_GT(argc, 0);
-
-  bootloader_message boot = {};
-  std::string err;
-  if (!read_bootloader_message(&boot, &err)) {
-    LOG(ERROR) << err;
-    // If fails, leave a zeroed bootloader_message.
-    boot = {};
-  }
-  stage = std::string(boot.stage);
-
-  if (boot.command[0] != 0) {
-    std::string boot_command = std::string(boot.command, sizeof(boot.command));
-    LOG(INFO) << "Boot command: " << boot_command;
-  }
-
-  if (boot.status[0] != 0) {
-    std::string boot_status = std::string(boot.status, sizeof(boot.status));
-    LOG(INFO) << "Boot status: " << boot_status;
-  }
-
-  std::vector<std::string> args(argv, argv + argc);
-
-  // --- if arguments weren't supplied, look in the bootloader control block
-  if (args.size() == 1) {
-    boot.recovery[sizeof(boot.recovery) - 1] = '\0';  // Ensure termination
-    std::string boot_recovery(boot.recovery);
-    std::vector<std::string> tokens = android::base::Split(boot_recovery, "\n");
-    if (!tokens.empty() && tokens[0] == "recovery") {
-      for (auto it = tokens.begin() + 1; it != tokens.end(); it++) {
-        // Skip empty and '\0'-filled tokens.
-        if (!it->empty() && (*it)[0] != '\0') args.push_back(std::move(*it));
-      }
-      LOG(INFO) << "Got " << args.size() << " arguments from boot message";
-    } else if (boot.recovery[0] != 0) {
-      LOG(ERROR) << "Bad boot message: \"" << boot_recovery << "\"";
-    }
-  }
-
-  // --- if that doesn't work, try the command file (if we have /cache).
-  if (args.size() == 1 && has_cache) {
-    std::string content;
-    if (ensure_path_mounted(COMMAND_FILE) == 0 &&
-        android::base::ReadFileToString(COMMAND_FILE, &content)) {
-      std::vector<std::string> tokens = android::base::Split(content, "\n");
-      // All the arguments in COMMAND_FILE are needed (unlike the BCB message,
-      // COMMAND_FILE doesn't use filename as the first argument).
-      for (auto it = tokens.begin(); it != tokens.end(); it++) {
-        // Skip empty and '\0'-filled tokens.
-        if (!it->empty() && (*it)[0] != '\0') args.push_back(std::move(*it));
-      }
-      LOG(INFO) << "Got " << args.size() << " arguments from " << COMMAND_FILE;
-    }
-  }
-
-  // Write the arguments (excluding the filename in args[0]) back into the
-  // bootloader control block. So the device will always boot into recovery to
-  // finish the pending work, until finish_recovery() is called.
-  std::vector<std::string> options(args.cbegin() + 1, args.cend());
-  if (!update_bootloader_message(options, &err)) {
-    LOG(ERROR) << "Failed to set BCB message: " << err;
-  }
-
-  return args;
-}
-
 // Set the BCB to reboot back into recovery (it won't resume the install from
 // sdcard though).
 static void set_sdcard_update_bootloader_message() {
@@ -921,21 +840,6 @@ static void print_property(const char* key, const char* name, void* /* cookie */
   printf("%s=%s\n", key, name);
 }
 
-static std::string load_locale_from_cache() {
-    if (ensure_path_mounted(LOCALE_FILE) != 0) {
-        LOG(ERROR) << "Can't mount " << LOCALE_FILE;
-        return "";
-    }
-
-    std::string content;
-    if (!android::base::ReadFileToString(LOCALE_FILE, &content)) {
-        PLOG(ERROR) << "Can't read " << LOCALE_FILE;
-        return "";
-    }
-
-    return android::base::Trim(content);
-}
-
 void ui_print(const char* format, ...) {
     std::string buffer;
     va_list ap;
@@ -1079,15 +983,7 @@ static void log_failure_code(ErrorCode code, const std::string& update_package)
   LOG(INFO) << log_content;
 }
 
-int start_recovery(int argc, char** argv) {
-  time_t start = time(nullptr);
-
-  printf("Starting recovery (pid %d) on %s", getpid(), ctime(&start));
-
-  load_volume_table();
-  has_cache = volume_for_mount_point(CACHE_ROOT) != nullptr;
-
-  std::vector<std::string> args = get_args(argc, argv);
+Device::BuiltinAction start_recovery(Device* device, const std::vector<std::string>& args) {
   std::vector<char*> args_to_parse(args.size());
   std::transform(args.cbegin(), args.cend(), args_to_parse.begin(),
                  [](const std::string& arg) { return const_cast<char*>(arg.c_str()); });
@@ -1117,7 +1013,6 @@ int start_recovery(int argc, char** argv) {
   bool should_wipe_cache = false;
   bool should_wipe_ab = false;
   size_t wipe_package_size = 0;
-  bool show_text = false;
   bool sideload = false;
   bool sideload_auto_reboot = false;
   bool just_exit = false;
@@ -1132,7 +1027,7 @@ int start_recovery(int argc, char** argv) {
                             &option_index)) != -1) {
     switch (arg) {
       case 't':
-        show_text = true;
+        // Handled in recovery_main.cpp
         break;
       case 'x':
         just_exit = true;
@@ -1140,7 +1035,7 @@ int start_recovery(int argc, char** argv) {
       case 0: {
         std::string option = OPTIONS[option_index].name;
         if (option == "locale") {
-          locale = optarg;
+          // Handled in recovery_main.cpp
         } else if (option == "prompt_and_wipe_data") {
           should_prompt_and_wipe_data = true;
         } else if (option == "reason") {
@@ -1174,38 +1069,11 @@ int start_recovery(int argc, char** argv) {
         continue;
     }
   }
+  optind = 1;
 
-  if (locale.empty()) {
-    if (has_cache) {
-      locale = load_locale_from_cache();
-    }
-
-    if (locale.empty()) {
-      static constexpr const char* DEFAULT_LOCALE = "en-US";
-      locale = DEFAULT_LOCALE;
-    }
-  }
-
-  printf("locale is [%s]\n", locale.c_str());
   printf("stage is [%s]\n", stage.c_str());
   printf("reason is [%s]\n", reason);
 
-  Device* device = make_device();
-  if (android::base::GetBoolProperty("ro.boot.quiescent", false)) {
-    printf("Quiescent recovery mode.\n");
-    device->ResetUI(new StubRecoveryUI());
-  } else {
-    if (!device->GetUI()->Init(locale)) {
-      printf("Failed to initialize UI; using stub UI instead.\n");
-      device->ResetUI(new StubRecoveryUI());
-    }
-  }
-  ui = device->GetUI();
-
-  if (!has_cache) {
-    device->RemoveMenuItemForAction(Device::WIPE_CACHE);
-  }
-
   // Set background string to "installing security update" for security update,
   // otherwise set it to "installing system update".
   ui->SetSystemUpdateText(security_update);
@@ -1215,15 +1083,6 @@ int start_recovery(int argc, char** argv) {
     ui->SetStage(st_cur, st_max);
   }
 
-  ui->SetBackground(RecoveryUI::NONE);
-  if (show_text) ui->ShowText(true);
-
-  sehandle = selinux_android_file_context_handle();
-  selinux_android_set_sehandle(sehandle);
-  if (!sehandle) {
-    ui->Print("Warning: No file_contexts\n");
-  }
-
   device->StartRecovery();
 
   printf("Command:");
@@ -1373,25 +1232,5 @@ int start_recovery(int argc, char** argv) {
   // Save logs and clean up before rebooting or shutting down.
   finish_recovery();
 
-  switch (after) {
-    case Device::SHUTDOWN:
-      ui->Print("Shutting down...\n");
-      android::base::SetProperty(ANDROID_RB_PROPERTY, "shutdown,");
-      break;
-
-    case Device::REBOOT_BOOTLOADER:
-      ui->Print("Rebooting to bootloader...\n");
-      android::base::SetProperty(ANDROID_RB_PROPERTY, "reboot,bootloader");
-      break;
-
-    default:
-      ui->Print("Rebooting...\n");
-      reboot("reboot,");
-      break;
-  }
-  while (true) {
-    pause();
-  }
-  // Should be unreachable.
-  return EXIT_SUCCESS;
+  return after;
 }
diff --git a/private/recovery.h b/recovery.h
index 5b2ca4b3f..00e22daa6 100644
--- a/private/recovery.h
+++ b/recovery.h
@@ -16,4 +16,9 @@
 
 #pragma once
 
-int start_recovery(int argc, char** argv);
+#include <string>
+#include <vector>
+
+#include "device.h"
+
+Device::BuiltinAction start_recovery(Device* device, const std::vector<std::string>& args);
diff --git a/recovery_main.cpp b/recovery_main.cpp
index 3147511ee..5e82c6c1a 100644
--- a/recovery_main.cpp
+++ b/recovery_main.cpp
@@ -14,22 +14,57 @@
  * limitations under the License.
  */
 
+#include <errno.h>
+#include <fcntl.h>
+#include <getopt.h>
+#include <inttypes.h>
+#include <limits.h>
+#include <linux/fs.h>
+#include <stdarg.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <time.h>
 #include <unistd.h>
 
-#include <chrono>
+#include <algorithm>
+#include <string>
+#include <vector>
 
+#include <android-base/file.h>
 #include <android-base/logging.h>
+#include <android-base/properties.h>
+#include <android-base/strings.h>
+#include <bootloader_message/bootloader_message.h>
+#include <cutils/android_reboot.h>
 #include <private/android_logger.h> /* private pmsg functions */
+#include <selinux/android.h>
+#include <selinux/label.h>
+#include <selinux/selinux.h>
 
 #include "common.h"
+#include "device.h"
 #include "logging.h"
 #include "minadbd/minadbd.h"
 #include "otautil/paths.h"
-#include "private/recovery.h"
+#include "otautil/sysutil.h"
+#include "recovery.h"
+#include "roots.h"
+#include "stub_ui.h"
 #include "ui.h"
 
+static constexpr const char* COMMAND_FILE = "/cache/recovery/command";
+static constexpr const char* LOCALE_FILE = "/cache/recovery/last_locale";
+
+static constexpr const char* CACHE_ROOT = "/cache";
+
+bool has_cache = false;
+
+RecoveryUI* ui = nullptr;
+struct selabel_handle* sehandle;
+
 static void UiLogger(android::base::LogId /* id */, android::base::LogSeverity severity,
                      const char* /* tag */, const char* /* file */, unsigned int /* line */,
                      const char* message) {
@@ -41,6 +76,92 @@ static void UiLogger(android::base::LogId /* id */, android::base::LogSeverity s
   }
 }
 
+// command line args come from, in decreasing precedence:
+//   - the actual command line
+//   - the bootloader control block (one per line, after "recovery")
+//   - the contents of COMMAND_FILE (one per line)
+static std::vector<std::string> get_args(const int argc, char** const argv) {
+  CHECK_GT(argc, 0);
+
+  bootloader_message boot = {};
+  std::string err;
+  if (!read_bootloader_message(&boot, &err)) {
+    LOG(ERROR) << err;
+    // If fails, leave a zeroed bootloader_message.
+    boot = {};
+  }
+  stage = std::string(boot.stage);
+
+  if (boot.command[0] != 0) {
+    std::string boot_command = std::string(boot.command, sizeof(boot.command));
+    LOG(INFO) << "Boot command: " << boot_command;
+  }
+
+  if (boot.status[0] != 0) {
+    std::string boot_status = std::string(boot.status, sizeof(boot.status));
+    LOG(INFO) << "Boot status: " << boot_status;
+  }
+
+  std::vector<std::string> args(argv, argv + argc);
+
+  // --- if arguments weren't supplied, look in the bootloader control block
+  if (args.size() == 1) {
+    boot.recovery[sizeof(boot.recovery) - 1] = '\0';  // Ensure termination
+    std::string boot_recovery(boot.recovery);
+    std::vector<std::string> tokens = android::base::Split(boot_recovery, "\n");
+    if (!tokens.empty() && tokens[0] == "recovery") {
+      for (auto it = tokens.begin() + 1; it != tokens.end(); it++) {
+        // Skip empty and '\0'-filled tokens.
+        if (!it->empty() && (*it)[0] != '\0') args.push_back(std::move(*it));
+      }
+      LOG(INFO) << "Got " << args.size() << " arguments from boot message";
+    } else if (boot.recovery[0] != 0) {
+      LOG(ERROR) << "Bad boot message: \"" << boot_recovery << "\"";
+    }
+  }
+
+  // --- if that doesn't work, try the command file (if we have /cache).
+  if (args.size() == 1 && has_cache) {
+    std::string content;
+    if (ensure_path_mounted(COMMAND_FILE) == 0 &&
+        android::base::ReadFileToString(COMMAND_FILE, &content)) {
+      std::vector<std::string> tokens = android::base::Split(content, "\n");
+      // All the arguments in COMMAND_FILE are needed (unlike the BCB message,
+      // COMMAND_FILE doesn't use filename as the first argument).
+      for (auto it = tokens.begin(); it != tokens.end(); it++) {
+        // Skip empty and '\0'-filled tokens.
+        if (!it->empty() && (*it)[0] != '\0') args.push_back(std::move(*it));
+      }
+      LOG(INFO) << "Got " << args.size() << " arguments from " << COMMAND_FILE;
+    }
+  }
+
+  // Write the arguments (excluding the filename in args[0]) back into the
+  // bootloader control block. So the device will always boot into recovery to
+  // finish the pending work, until finish_recovery() is called.
+  std::vector<std::string> options(args.cbegin() + 1, args.cend());
+  if (!update_bootloader_message(options, &err)) {
+    LOG(ERROR) << "Failed to set BCB message: " << err;
+  }
+
+  return args;
+}
+
+static std::string load_locale_from_cache() {
+  if (ensure_path_mounted(LOCALE_FILE) != 0) {
+    LOG(ERROR) << "Can't mount " << LOCALE_FILE;
+    return "";
+  }
+
+  std::string content;
+  if (!android::base::ReadFileToString(LOCALE_FILE, &content)) {
+    PLOG(ERROR) << "Can't read " << LOCALE_FILE;
+    return "";
+  }
+
+  return android::base::Trim(content);
+}
+
 static void redirect_stdio(const char* filename) {
   int pipefd[2];
   if (pipe(pipefd) == -1) {
@@ -154,9 +275,109 @@ int main(int argc, char** argv) {
     return 0;
   }
 
+  time_t start = time(nullptr);
+
   // redirect_stdio should be called only in non-sideload mode. Otherwise we may have two logger
   // instances with different timestamps.
   redirect_stdio(Paths::Get().temporary_log_file().c_str());
 
-  return start_recovery(argc, argv);
+  printf("Starting recovery (pid %d) on %s", getpid(), ctime(&start));
+
+  load_volume_table();
+  has_cache = volume_for_mount_point(CACHE_ROOT) != nullptr;
+
+  std::vector<std::string> args = get_args(argc, argv);
+  std::vector<char*> args_to_parse(args.size());
+  std::transform(args.cbegin(), args.cend(), args_to_parse.begin(),
+                 [](const std::string& arg) { return const_cast<char*>(arg.c_str()); });
+
+  static constexpr struct option OPTIONS[] = {
+    { "locale", required_argument, nullptr, 0 },
+    { "show_text", no_argument, nullptr, 't' },
+    { nullptr, 0, nullptr, 0 },
+  };
+
+  bool show_text = false;
+  std::string locale;
+
+  int arg;
+  int option_index;
+  while ((arg = getopt_long(args_to_parse.size(), args_to_parse.data(), "", OPTIONS,
+                            &option_index)) != -1) {
+    switch (arg) {
+      case 't':
+        show_text = true;
+        break;
+      case 0: {
+        std::string option = OPTIONS[option_index].name;
+        if (option == "locale") {
+          locale = optarg;
+        }
+        break;
+      }
+    }
+  }
+  optind = 1;
+
+  if (locale.empty()) {
+    if (has_cache) {
+      locale = load_locale_from_cache();
+    }
+
+    if (locale.empty()) {
+      static constexpr const char* DEFAULT_LOCALE = "en-US";
+      locale = DEFAULT_LOCALE;
+    }
+  }
+
+  printf("locale is [%s]\n", locale.c_str());
+
+  Device* device = make_device();
+  if (android::base::GetBoolProperty("ro.boot.quiescent", false)) {
+    printf("Quiescent recovery mode.\n");
+    device->ResetUI(new StubRecoveryUI());
+  } else {
+    if (!device->GetUI()->Init(locale)) {
+      printf("Failed to initialize UI; using stub UI instead.\n");
+      device->ResetUI(new StubRecoveryUI());
+    }
+  }
+  ui = device->GetUI();
+
+  if (!has_cache) {
+    device->RemoveMenuItemForAction(Device::WIPE_CACHE);
+  }
+
+  ui->SetBackground(RecoveryUI::NONE);
+  if (show_text) ui->ShowText(true);
+
+  sehandle = selinux_android_file_context_handle();
+  selinux_android_set_sehandle(sehandle);
+  if (!sehandle) {
+    ui->Print("Warning: No file_contexts\n");
+  }
+
+  Device::BuiltinAction after = start_recovery(device, args);
+
+  switch (after) {
+    case Device::SHUTDOWN:
+      ui->Print("Shutting down...\n");
+      android::base::SetProperty(ANDROID_RB_PROPERTY, "shutdown,");
+      break;
+
+    case Device::REBOOT_BOOTLOADER:
+      ui->Print("Rebooting to bootloader...\n");
+      android::base::SetProperty(ANDROID_RB_PROPERTY, "reboot,bootloader");
+      break;
+
+    default:
+      ui->Print("Rebooting...\n");
+      reboot("reboot,");
+      break;
+  }
+  while (true) {
+    pause();
+  }
+  // Should be unreachable.
+  return EXIT_SUCCESS;
 }
diff --git a/tests/unit/screen_ui_test.cpp b/tests/unit/screen_ui_test.cpp
index 03e23ca42..269222faa 100644
--- a/tests/unit/screen_ui_test.cpp
+++ b/tests/unit/screen_ui_test.cpp
@@ -15,6 +15,7 @@
  */
 
 #include <stddef.h>
+#include <stdio.h>
 
 #include <functional>
 #include <map>
@@ -23,6 +24,8 @@
 #include <vector>
 
 #include <android-base/logging.h>
+#include <android-base/stringprintf.h>
+#include <android-base/test_utils.h>
 #include <gtest/gtest.h>
 
 #include "common/test_constants.h"
@@ -224,6 +227,19 @@ class TestableScreenRecoveryUI : public ScreenRecoveryUI {
 
   int KeyHandler(int key, bool visible) const;
 
+  // The following functions expose the protected members for test purpose.
+  void RunLoadAnimation() {
+    LoadAnimation();
+  }
+
+  size_t GetLoopFrames() const {
+    return loop_frames;
+  }
+
+  size_t GetIntroFrames() const {
+    return intro_frames;
+  }
+
   bool GetRtlLocale() const {
     return rtl_locale_;
   }
@@ -260,14 +276,15 @@ class ScreenRecoveryUITest : public ::testing::Test {
   void SetUp() override {
     ui_ = std::make_unique<TestableScreenRecoveryUI>();
 
-    std::string testdata_dir = from_testdata_base("");
-    Paths::Get().set_resource_dir(testdata_dir);
-    res_set_resource_dir(testdata_dir);
+    testdata_dir_ = from_testdata_base("");
+    Paths::Get().set_resource_dir(testdata_dir_);
+    res_set_resource_dir(testdata_dir_);
 
     ASSERT_TRUE(ui_->Init(kTestLocale));
   }
 
   std::unique_ptr<TestableScreenRecoveryUI> ui_;
+  std::string testdata_dir_;
 };
 
 TEST_F(ScreenRecoveryUITest, Init) {
@@ -352,3 +369,36 @@ TEST_F(ScreenRecoveryUITest, ShowMenu_TimedOut_TextWasEverVisible) {
                               std::bind(&TestableScreenRecoveryUI::KeyHandler, ui_.get(),
                                         std::placeholders::_1, std::placeholders::_2)));
 }
+
+TEST_F(ScreenRecoveryUITest, LoadAnimation) {
+  // Make a few copies of loop00000.png from testdata.
+  std::string image_data;
+  ASSERT_TRUE(android::base::ReadFileToString(testdata_dir_ + "/loop00000.png", &image_data));
+
+  std::vector<std::string> tempfiles;
+  TemporaryDir resource_dir;
+  for (const auto& name : { "00002", "00100", "00050" }) {
+    tempfiles.push_back(android::base::StringPrintf("%s/loop%s.png", resource_dir.path, name));
+    ASSERT_TRUE(android::base::WriteStringToFile(image_data, tempfiles.back()));
+  }
+  for (const auto& name : { "00", "01" }) {
+    tempfiles.push_back(android::base::StringPrintf("%s/intro%s.png", resource_dir.path, name));
+    ASSERT_TRUE(android::base::WriteStringToFile(image_data, tempfiles.back()));
+  }
+  Paths::Get().set_resource_dir(resource_dir.path);
+
+  ui_->RunLoadAnimation();
+
+  ASSERT_EQ(2u, ui_->GetIntroFrames());
+  ASSERT_EQ(3u, ui_->GetLoopFrames());
+
+  for (const auto& name : tempfiles) {
+    ASSERT_EQ(0, unlink(name.c_str()));
+  }
+}
+
+TEST_F(ScreenRecoveryUITest, LoadAnimation_MissingAnimation) {
+  TemporaryDir resource_dir;
+  Paths::Get().set_resource_dir(resource_dir.path);
+  ASSERT_EXIT(ui_->RunLoadAnimation(), ::testing::KilledBySignal(SIGABRT), "");
+}
diff --git a/updater/blockimg.cpp b/updater/blockimg.cpp
index 156a82939..236644e7f 100644
--- a/updater/blockimg.cpp
+++ b/updater/blockimg.cpp
@@ -132,8 +132,7 @@ static bool FsyncDir(const std::string& dirname) {
   return true;
 }
 
-// Update the last command index in the last_command_file if the current command writes to the
-// stash either explicitly or implicitly.
+// Update the last executed command index in the last_command_file.
 static bool UpdateLastCommandIndex(int command_index, const std::string& command_string) {
   const std::string& last_command_file = Paths::Get().last_command_file();
   std::string last_command_tmp = last_command_file + ".tmp";
@@ -1161,10 +1160,6 @@ static int LoadSrcTgtVersion3(CommandParameters& params, RangeSet& tgt, size_t*
         return -1;
       }
 
-      if (!UpdateLastCommandIndex(params.cmdindex, params.cmdline)) {
-        LOG(WARNING) << "Failed to update the last command file.";
-      }
-
       params.stashed += *src_blocks;
       // Can be deleted when the write has completed.
       if (!stash_exists) {
@@ -1275,10 +1270,6 @@ static int PerformCommandStash(CommandParameters& params) {
   LOG(INFO) << "stashing " << blocks << " blocks to " << id;
   int result = WriteStash(params.stashbase, id, blocks, params.buffer, false, nullptr);
   if (result == 0) {
-    if (!UpdateLastCommandIndex(params.cmdindex, params.cmdline)) {
-      LOG(WARNING) << "Failed to update the last command file.";
-    }
-
     params.stashed += blocks;
   }
   return result;
@@ -1701,7 +1692,7 @@ static Value* PerformBlockImageUpdate(const char* name, State* state,
   params.createdstash = res;
 
   // When performing an update, save the index and cmdline of the current command into
-  // the last_command_file if this command writes to the stash either explicitly of implicitly.
+  // the last_command_file.
   // Upon resuming an update, read the saved index first; then
   //   1. In verification mode, check if the 'move' or 'diff' commands before the saved index has
   //      the expected target blocks already. If not, these commands cannot be skipped and we need
@@ -1797,6 +1788,11 @@ static Value* PerformBlockImageUpdate(const char* name, State* state,
         PLOG(ERROR) << "fsync failed";
         goto pbiudone;
       }
+
+      if (!UpdateLastCommandIndex(params.cmdindex, params.cmdline)) {
+        LOG(WARNING) << "Failed to update the last command file.";
+      }
+
       fprintf(cmd_pipe, "set_progress %.4f\n", static_cast<double>(params.written) / total_blocks);
       fflush(cmd_pipe);
     }
diff --git a/updater_sample/README.md b/updater_sample/README.md
index 95e57dbe9..c68c07caf 100644
--- a/updater_sample/README.md
+++ b/updater_sample/README.md
@@ -44,6 +44,10 @@ saved uncompressed (`ZIP_STORED`), so that their data can be downloaded directly
 with the offset and length. As `payload.bin` itself is already in compressed
 format, the size penalty is marginal.
 
+if `ab_config.force_switch_slot` set true device will boot to the
+updated partition on next reboot; otherwise button "Switch Slot" will
+become active, and user can manually set updated partition as the active slot.
+
 Config files can be generated using `tools/gen_update_config.py`.
 Running `./tools/gen_update_config.py --help` shows usage of the script.
 
@@ -85,8 +89,8 @@ which HTTP headers are supported.
 - [x] Add stop/reset the update
 - [x] Add demo for passing HTTP headers to `UpdateEngine#applyPayload`
 - [x] [Package compatibility check](https://source.android.com/devices/architecture/vintf/match-rules)
+- [x] Deferred switch slot demo
 - [ ] Add tests for `MainActivity`
-- [ ] Change partition demo
 - [ ] Verify system partition checksum for package
 - [ ] Add non-A/B updates demo
 
diff --git a/updater_sample/res/layout/activity_main.xml b/updater_sample/res/layout/activity_main.xml
index 7a12d3474..d9e56b4b3 100644
--- a/updater_sample/res/layout/activity_main.xml
+++ b/updater_sample/res/layout/activity_main.xml
@@ -178,6 +178,23 @@
                     android:text="Reset" />
             </LinearLayout>
 
+            <TextView
+                android:id="@+id/textViewUpdateInfo"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:layout_marginTop="14dp"
+                android:textColor="#777"
+                android:textSize="10sp"
+                android:textStyle="italic"
+                android:text="@string/finish_update_info" />
+
+            <Button
+                android:id="@+id/buttonSwitchSlot"
+                android:layout_width="wrap_content"
+                android:layout_height="match_parent"
+                android:onClick="onSwitchSlotClick"
+                android:text="@string/switch_slot" />
+
         </LinearLayout>
 
     </ScrollView>
diff --git a/updater_sample/res/raw/sample.json b/updater_sample/res/raw/sample.json
index 46fbfa33e..f188c233b 100644
--- a/updater_sample/res/raw/sample.json
+++ b/updater_sample/res/raw/sample.json
@@ -20,5 +20,10 @@
             }
         ],
         "authorization": "Basic my-secret-token"
+    },
+    "ab_config": {
+        "__": "A/B (seamless) update configurations",
+        "__force_switch_slot": "if set true device will boot to a new slot, otherwise user manually switches slot on the screen",
+        "force_switch_slot": false
     }
 }
diff --git a/updater_sample/res/values/strings.xml b/updater_sample/res/values/strings.xml
index 2b671ee5d..db4a5dc67 100644
--- a/updater_sample/res/values/strings.xml
+++ b/updater_sample/res/values/strings.xml
@@ -18,4 +18,6 @@
     <string name="action_reload">Reload</string>
     <string name="unknown">Unknown</string>
     <string name="close">CLOSE</string>
+    <string name="switch_slot">Switch Slot</string>
+    <string name="finish_update_info">To finish the update press the button below</string>
 </resources>
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java
index 9bdd8b9e8..db99f7c74 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java
@@ -71,7 +71,7 @@ public class UpdateConfig implements Parcelable {
             JSONObject meta = o.getJSONObject("ab_streaming_metadata");
             JSONArray propertyFilesJson = meta.getJSONArray("property_files");
             PackageFile[] propertyFiles =
-                new PackageFile[propertyFilesJson.length()];
+                    new PackageFile[propertyFilesJson.length()];
             for (int i = 0; i < propertyFilesJson.length(); i++) {
                 JSONObject p = propertyFilesJson.getJSONObject(i);
                 propertyFiles[i] = new PackageFile(
@@ -87,6 +87,12 @@ public class UpdateConfig implements Parcelable {
                     propertyFiles,
                     authorization);
         }
+
+        // TODO: parse only for A/B updates when non-A/B is implemented
+        JSONObject ab = o.getJSONObject("ab_config");
+        boolean forceSwitchSlot = ab.getBoolean("force_switch_slot");
+        c.mAbConfig = new AbConfig(forceSwitchSlot);
+
         c.mRawJson = json;
         return c;
     }
@@ -109,6 +115,9 @@ public class UpdateConfig implements Parcelable {
     /** metadata is required only for streaming update */
     private StreamingMetadata mAbStreamingMetadata;
 
+    /** A/B update configurations */
+    private AbConfig mAbConfig;
+
     private String mRawJson;
 
     protected UpdateConfig() {
@@ -119,6 +128,7 @@ public class UpdateConfig implements Parcelable {
         this.mUrl = in.readString();
         this.mAbInstallType = in.readInt();
         this.mAbStreamingMetadata = (StreamingMetadata) in.readSerializable();
+        this.mAbConfig = (AbConfig) in.readSerializable();
         this.mRawJson = in.readString();
     }
 
@@ -148,6 +158,10 @@ public class UpdateConfig implements Parcelable {
         return mAbStreamingMetadata;
     }
 
+    public AbConfig getAbConfig() {
+        return mAbConfig;
+    }
+
     /**
      * @return File object for given url
      */
@@ -172,6 +186,7 @@ public class UpdateConfig implements Parcelable {
         dest.writeString(mUrl);
         dest.writeInt(mAbInstallType);
         dest.writeSerializable(mAbStreamingMetadata);
+        dest.writeSerializable(mAbConfig);
         dest.writeString(mRawJson);
     }
 
@@ -185,9 +200,11 @@ public class UpdateConfig implements Parcelable {
         /** defines beginning of update data in archive */
         private PackageFile[] mPropertyFiles;
 
-        /** SystemUpdaterSample receives the authorization token from the OTA server, in addition
+        /**
+         * SystemUpdaterSample receives the authorization token from the OTA server, in addition
          * to the package URL. It passes on the info to update_engine, so that the latter can
-         * fetch the data from the package server directly with the token. */
+         * fetch the data from the package server directly with the token.
+         */
         private String mAuthorization;
 
         public StreamingMetadata(PackageFile[] propertyFiles, String authorization) {
@@ -239,4 +256,27 @@ public class UpdateConfig implements Parcelable {
         }
     }
 
-}
+    /**
+     * A/B (seamless) update configurations.
+     */
+    public static class AbConfig implements Serializable {
+
+        private static final long serialVersionUID = 31044L;
+
+        /**
+         * if set true device will boot to new slot, otherwise user manually
+         * switches slot on the screen.
+         */
+        private boolean mForceSwitchSlot;
+
+        public AbConfig(boolean forceSwitchSlot) {
+            this.mForceSwitchSlot = forceSwitchSlot;
+        }
+
+        public boolean getForceSwitchSlot() {
+            return mForceSwitchSlot;
+        }
+
+    }
+
+}
\ No newline at end of file
diff --git a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
index 170825635..c5a7f9556 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -18,6 +18,7 @@ package com.example.android.systemupdatersample.ui;
 
 import android.app.Activity;
 import android.app.AlertDialog;
+import android.graphics.Color;
 import android.os.Build;
 import android.os.Bundle;
 import android.os.UpdateEngine;
@@ -38,11 +39,13 @@ import com.example.android.systemupdatersample.services.PrepareStreamingService;
 import com.example.android.systemupdatersample.util.PayloadSpecs;
 import com.example.android.systemupdatersample.util.UpdateConfigs;
 import com.example.android.systemupdatersample.util.UpdateEngineErrorCodes;
+import com.example.android.systemupdatersample.util.UpdateEngineProperties;
 import com.example.android.systemupdatersample.util.UpdateEngineStatuses;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -66,10 +69,14 @@ public class MainActivity extends Activity {
     private ProgressBar mProgressBar;
     private TextView mTextViewStatus;
     private TextView mTextViewCompletion;
+    private TextView mTextViewUpdateInfo;
+    private Button mButtonSwitchSlot;
 
     private List<UpdateConfig> mConfigs;
     private AtomicInteger mUpdateEngineStatus =
             new AtomicInteger(UpdateEngine.UpdateStatusConstants.IDLE);
+    private PayloadSpec mLastPayloadSpec;
+    private AtomicBoolean mManualSwitchSlotRequired = new AtomicBoolean(true);
 
     /**
      * Listen to {@code update_engine} events.
@@ -93,6 +100,8 @@ public class MainActivity extends Activity {
         this.mProgressBar = findViewById(R.id.progressBar);
         this.mTextViewStatus = findViewById(R.id.textViewStatus);
         this.mTextViewCompletion = findViewById(R.id.textViewCompletion);
+        this.mTextViewUpdateInfo = findViewById(R.id.textViewUpdateInfo);
+        this.mButtonSwitchSlot = findViewById(R.id.buttonSwitchSlot);
 
         this.mTextViewConfigsDirHint.setText(UpdateConfigs.getConfigsRoot(this));
 
@@ -173,6 +182,13 @@ public class MainActivity extends Activity {
     }
 
     /**
+     * switch slot button clicked
+     */
+    public void onSwitchSlotClick(View view) {
+        setSwitchSlotOnReboot();
+    }
+
+    /**
      * Invoked when anything changes. The value of {@code status} will
      * be one of the values from {@link UpdateEngine.UpdateStatusConstants},
      * and {@code percent} will be from {@code 0.0} to {@code 1.0}.
@@ -185,16 +201,16 @@ public class MainActivity extends Activity {
                 Log.e("UpdateEngine", "StatusUpdate - status="
                         + UpdateEngineStatuses.getStatusText(status)
                         + "/" + status);
-                setUiStatus(status);
                 Toast.makeText(this, "Update Status changed", Toast.LENGTH_LONG)
                         .show();
-                if (status != UpdateEngine.UpdateStatusConstants.IDLE) {
-                    Log.d(TAG, "status changed, setting ui to updating mode");
-                    uiSetUpdating();
-                } else {
+                if (status == UpdateEngine.UpdateStatusConstants.IDLE) {
                     Log.d(TAG, "status changed, resetting ui");
                     uiReset();
+                } else {
+                    Log.d(TAG, "status changed, setting ui to updating mode");
+                    uiSetUpdating();
                 }
+                setUiStatus(status);
             });
         }
     }
@@ -215,6 +231,13 @@ public class MainActivity extends Activity {
                     + " " + state);
             Toast.makeText(this, "Update completed", Toast.LENGTH_LONG).show();
             setUiCompletion(errorCode);
+            if (errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
+                // if update was successfully applied.
+                if (mManualSwitchSlotRequired.get()) {
+                    // Show "Switch Slot" button.
+                    uiShowSwitchSlotInfo();
+                }
+            }
         });
     }
 
@@ -231,6 +254,7 @@ public class MainActivity extends Activity {
         mProgressBar.setVisibility(ProgressBar.INVISIBLE);
         mTextViewStatus.setText(R.string.unknown);
         mTextViewCompletion.setText(R.string.unknown);
+        uiHideSwitchSlotInfo();
     }
 
     /** sets ui updating mode */
@@ -245,6 +269,16 @@ public class MainActivity extends Activity {
         mProgressBar.setVisibility(ProgressBar.VISIBLE);
     }
 
+    private void uiShowSwitchSlotInfo() {
+        mButtonSwitchSlot.setEnabled(true);
+        mTextViewUpdateInfo.setTextColor(Color.parseColor("#777777"));
+    }
+
+    private void uiHideSwitchSlotInfo() {
+        mTextViewUpdateInfo.setTextColor(Color.parseColor("#AAAAAA"));
+        mButtonSwitchSlot.setEnabled(false);
+    }
+
     /**
      * loads json configurations from configs dir that is defined in {@link UpdateConfigs}.
      */
@@ -290,6 +324,17 @@ public class MainActivity extends Activity {
      * Applies the given update
      */
     private void applyUpdate(final UpdateConfig config) {
+        List<String> extraProperties = new ArrayList<>();
+
+        if (!config.getAbConfig().getForceSwitchSlot()) {
+            // Disable switch slot on reboot, which is enabled by default.
+            // User will enable it manually by clicking "Switch Slot" button on the screen.
+            extraProperties.add(UpdateEngineProperties.PROPERTY_DISABLE_SWITCH_SLOT_ON_REBOOT);
+            mManualSwitchSlotRequired.set(true);
+        } else {
+            mManualSwitchSlotRequired.set(false);
+        }
+
         if (config.getInstallType() == UpdateConfig.AB_INSTALL_TYPE_NON_STREAMING) {
             PayloadSpec payload;
             try {
@@ -300,12 +345,11 @@ public class MainActivity extends Activity {
                         .show();
                 return;
             }
-            updateEngineApplyPayload(payload, null);
+            updateEngineApplyPayload(payload, extraProperties);
         } else {
             Log.d(TAG, "Starting PrepareStreamingService");
             PrepareStreamingService.startService(this, config, (code, payloadSpec) -> {
                 if (code == PrepareStreamingService.RESULT_CODE_SUCCESS) {
-                    List<String> extraProperties = new ArrayList<>();
                     extraProperties.add("USER_AGENT=" + HTTP_USER_AGENT);
                     config.getStreamingMetadata()
                             .getAuthorization()
@@ -332,6 +376,8 @@ public class MainActivity extends Activity {
      * @param extraProperties additional properties to pass to {@link UpdateEngine#applyPayload}
      */
     private void updateEngineApplyPayload(PayloadSpec payloadSpec, List<String> extraProperties) {
+        mLastPayloadSpec = payloadSpec;
+
         ArrayList<String> properties = new ArrayList<>(payloadSpec.getProperties());
         if (extraProperties != null) {
             properties.addAll(extraProperties);
@@ -352,6 +398,29 @@ public class MainActivity extends Activity {
     }
 
     /**
+     * Sets the new slot that has the updated partitions as the active slot,
+     * which device will boot into next time.
+     * This method is only supposed to be called after the payload is applied.
+     *
+     * Invoking {@link UpdateEngine#applyPayload} with the same payload url, offset, size
+     * and payload metadata headers doesn't trigger new update. It can be used to just switch
+     * active A/B slot.
+     *
+     * {@link UpdateEngine#applyPayload} might take several seconds to finish, and it will
+     * invoke callbacks {@link this#onStatusUpdate} and {@link this#onPayloadApplicationComplete)}.
+     */
+    private void setSwitchSlotOnReboot() {
+        Log.d(TAG, "setSwitchSlotOnReboot invoked");
+        List<String> extraProperties = new ArrayList<>();
+        // PROPERTY_SKIP_POST_INSTALL should be passed on to skip post-installation hooks.
+        extraProperties.add(UpdateEngineProperties.PROPERTY_SKIP_POST_INSTALL);
+        // It sets property SWITCH_SLOT_ON_REBOOT=1 by default.
+        // HTTP headers are not required, UpdateEngine is not expected to stream payload.
+        updateEngineApplyPayload(mLastPayloadSpec, extraProperties);
+        uiHideSwitchSlotInfo();
+    }
+
+    /**
      * Requests update engine to stop any ongoing update. If an update has been applied,
      * leave it as is.
      */
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineProperties.java b/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineProperties.java
new file mode 100644
index 000000000..e368f14d2
--- /dev/null
+++ b/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineProperties.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.android.systemupdatersample.util;
+
+/**
+ * Utility class for properties that will be passed to {@code UpdateEngine#applyPayload}.
+ */
+public final class UpdateEngineProperties {
+
+    /**
+     * The property indicating that the update engine should not switch slot
+     * when the device reboots.
+     */
+    public static final String PROPERTY_DISABLE_SWITCH_SLOT_ON_REBOOT = "SWITCH_SLOT_ON_REBOOT=0";
+
+    /**
+     * The property to skip post-installation.
+     * https://source.android.com/devices/tech/ota/ab/#post-installation
+     */
+    public static final String PROPERTY_SKIP_POST_INSTALL = "RUN_POST_INSTALL=0";
+
+    private UpdateEngineProperties() {}
+}
diff --git a/updater_sample/tests/res/raw/update_config_stream_001.json b/updater_sample/tests/res/raw/update_config_stream_001.json
index 15127cf2c..be51b7c95 100644
--- a/updater_sample/tests/res/raw/update_config_stream_001.json
+++ b/updater_sample/tests/res/raw/update_config_stream_001.json
@@ -10,5 +10,8 @@
                 "size": 8
             }
         ]
+    },
+    "ab_config": {
+        "force_switch_slot": true
     }
 }
diff --git a/updater_sample/tests/res/raw/update_config_stream_002.json b/updater_sample/tests/res/raw/update_config_stream_002.json
index cf4469b1c..5d7874cdb 100644
--- a/updater_sample/tests/res/raw/update_config_stream_002.json
+++ b/updater_sample/tests/res/raw/update_config_stream_002.json
@@ -1,5 +1,8 @@
 {
     "__": "*** Generated using tools/gen_update_config.py ***",
+    "ab_config": {
+        "force_switch_slot": false
+    },
     "ab_install_type": "STREAMING",
     "ab_streaming_metadata": {
         "property_files": [
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
index 0975e76be..000f5663b 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
@@ -18,6 +18,7 @@ package com.example.android.systemupdatersample;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 
 import android.content.Context;
 import android.support.test.InstrumentationRegistry;
@@ -45,7 +46,8 @@ public class UpdateConfigTest {
 
     private static final String JSON_NON_STREAMING =
             "{\"name\": \"vip update\", \"url\": \"file:///builds/a.zip\", "
-            + " \"ab_install_type\": \"NON_STREAMING\"}";
+                    + " \"ab_install_type\": \"NON_STREAMING\","
+                    + " \"ab_config\": { \"force_switch_slot\": false } }";
 
     @Rule
     public final ExpectedException thrown = ExpectedException.none();
@@ -82,6 +84,7 @@ public class UpdateConfigTest {
                 config.getStreamingMetadata().getPropertyFiles()[0].getFilename());
         assertEquals(195, config.getStreamingMetadata().getPropertyFiles()[0].getOffset());
         assertEquals(8, config.getStreamingMetadata().getPropertyFiles()[0].getSize());
+        assertTrue(config.getAbConfig().getForceSwitchSlot());
     }
 
     @Test
@@ -94,7 +97,8 @@ public class UpdateConfigTest {
     @Test
     public void getUpdatePackageFile_throwsErrorIfNotAFile() throws Exception {
         String json = "{\"name\": \"upd\", \"url\": \"http://foo.bar\","
-                + " \"ab_install_type\": \"NON_STREAMING\"}";
+                + " \"ab_install_type\": \"NON_STREAMING\","
+                + " \"ab_config\": { \"force_switch_slot\": false } }";
         UpdateConfig config = UpdateConfig.fromJson(json);
         thrown.expect(RuntimeException.class);
         config.getUpdatePackageFile();
diff --git a/updater_sample/tools/gen_update_config.py b/updater_sample/tools/gen_update_config.py
index 4efa9f1c4..7fb64f7fc 100755
--- a/updater_sample/tools/gen_update_config.py
+++ b/updater_sample/tools/gen_update_config.py
@@ -46,10 +46,11 @@ class GenUpdateConfig(object):
     AB_INSTALL_TYPE_STREAMING = 'STREAMING'
     AB_INSTALL_TYPE_NON_STREAMING = 'NON_STREAMING'
 
-    def __init__(self, package, url, ab_install_type):
+    def __init__(self, package, url, ab_install_type, ab_force_switch_slot):
         self.package = package
         self.url = url
         self.ab_install_type = ab_install_type
+        self.ab_force_switch_slot = ab_force_switch_slot
         self.streaming_required = (
             # payload.bin and payload_properties.txt must exist.
             'payload.bin',
@@ -80,6 +81,9 @@ class GenUpdateConfig(object):
             'url': self.url,
             'ab_streaming_metadata': streaming_metadata,
             'ab_install_type': self.ab_install_type,
+            'ab_config': {
+                'force_switch_slot': self.ab_force_switch_slot,
+            }
         }
 
     def _gen_ab_streaming_metadata(self):
@@ -126,6 +130,11 @@ def main():  # pylint: disable=missing-docstring
                         default=GenUpdateConfig.AB_INSTALL_TYPE_NON_STREAMING,
                         choices=ab_install_type_choices,
                         help='A/B update installation type')
+    parser.add_argument('--ab_force_switch_slot',
+                        type=bool,
+                        default=False,
+                        help='if set true device will boot to a new slot, otherwise user manually '
+                             'switches slot on the screen')
     parser.add_argument('package',
                         type=str,
                         help='OTA package zip file')
@@ -144,7 +153,8 @@ def main():  # pylint: disable=missing-docstring
     gen = GenUpdateConfig(
         package=args.package,
         url=args.url,
-        ab_install_type=args.ab_install_type)
+        ab_install_type=args.ab_install_type,
+        ab_force_switch_slot=args.ab_force_switch_slot)
     gen.run()
     gen.write(args.out)
     print('Config is written to ' + args.out)
diff --git a/wear_ui.cpp b/wear_ui.cpp
index f157d3ca3..f4a839923 100644
--- a/wear_ui.cpp
+++ b/wear_ui.cpp
@@ -32,11 +32,6 @@ WearRecoveryUI::WearRecoveryUI()
       kMenuUnusableRows(RECOVERY_UI_MENU_UNUSABLE_ROWS) {
   // TODO: kMenuUnusableRows should be computed based on the lines in draw_screen_locked().
 
-  // TODO: The following three variables are likely not needed. The first two are detected
-  // automatically in ScreenRecoveryUI::LoadAnimation(), based on the actual files seen on device.
-  intro_frames = 22;
-  loop_frames = 60;
-
   touch_screen_allowed_ = true;
 }