diff --git a/client_library/client_impl.cc b/client_library/client_impl.cc
index 7b80025..247f8f0 100644
--- a/client_library/client_impl.cc
+++ b/client_library/client_impl.cc
@@ -40,10 +40,8 @@
                                            const string& in_omaha_url,
                                            bool at_user_request) {
   return proxy_->AttemptUpdateWithFlags(
-      in_app_version,
-      in_omaha_url,
-      (at_user_request) ? 0 : kAttemptUpdateFlagNonInteractive,
-      nullptr);
+      in_app_version, in_omaha_url,
+      (at_user_request) ? 0 : kAttemptUpdateFlagNonInteractive, nullptr);
 }
 
 bool UpdateEngineClientImpl::GetStatus(int64_t* out_last_checked_time,
@@ -52,13 +50,9 @@
                                        string* out_new_version,
                                        int64_t* out_new_size) {
   string status_as_string;
-  const bool success = proxy_->GetStatus(
-      out_last_checked_time,
-      out_progress,
-      &status_as_string,
-      out_new_version,
-      out_new_size,
-      nullptr);
+  const bool success =
+      proxy_->GetStatus(out_last_checked_time, out_progress, &status_as_string,
+                        out_new_version, out_new_size, nullptr);
   if (!success) {
     return false;
   }
@@ -70,7 +64,7 @@
   return proxy_->SetUpdateOverCellularPermission(allowed, nullptr);
 }
 
-bool UpdateEngineClientImpl::GetUpdateOverCellularPermission(bool *allowed) {
+bool UpdateEngineClientImpl::GetUpdateOverCellularPermission(bool* allowed) {
   return proxy_->GetUpdateOverCellularPermission(allowed, nullptr);
 }
 
@@ -108,26 +102,63 @@
   return proxy_->ResetStatus(nullptr);
 }
 
+void UpdateEngineClientImpl::StatusUpdateHandlerRegistered(
+    StatusUpdateHandler* handler, const std::string& interface,
+    const std::string& signal_name, bool success) {
+  if (!success) {
+    handler->IPCError("Could not connect to" + signal_name);
+    return;
+  }
+
+  int64_t last_checked_time;
+  double progress;
+  UpdateStatus update_status;
+  string new_version;
+  int64_t new_size;
+
+  if (GetStatus(&last_checked_time, &progress, &update_status, &new_version,
+                &new_size)) {
+    handler->HandleStatusUpdate(last_checked_time, progress, update_status,
+                                new_version, new_size);
+    return;
+  }
+
+  handler->IPCError("Could not query current status");
+}
+
+void UpdateEngineClientImpl::RunStatusUpdateHandler(
+    StatusUpdateHandler* h, int64_t last_checked_time, double progress,
+    const std::string& current_operation, const std::string& new_version,
+    int64_t new_size) {
+  UpdateStatus status;
+  StringToUpdateStatus(current_operation, &status);
+
+  h->HandleStatusUpdate(last_checked_time, progress, status, new_version,
+                        new_size);
+}
+
+void UpdateEngineClientImpl::RegisterStatusUpdateHandler(
+    StatusUpdateHandler* handler) {
+  proxy_->RegisterStatusUpdateSignalHandler(
+      base::Bind(&UpdateEngineClientImpl::RunStatusUpdateHandler,
+                 base::Unretained(this), base::Unretained(handler)),
+      base::Bind(&UpdateEngineClientImpl::StatusUpdateHandlerRegistered,
+                 base::Unretained(this), base::Unretained(handler)));
+}
+
 bool UpdateEngineClientImpl::SetTargetChannel(const string& in_target_channel,
                                               bool allow_powerwash) {
-  return proxy_->SetChannel(
-      in_target_channel,
-      allow_powerwash,
-      nullptr);
+  return proxy_->SetChannel(in_target_channel, allow_powerwash, nullptr);
 }
 
 bool UpdateEngineClientImpl::GetTargetChannel(string* out_channel) {
-  return proxy_->GetChannel(
-      false,  // Get the target channel.
-      out_channel,
-      nullptr);
+  return proxy_->GetChannel(false,  // Get the target channel.
+                            out_channel, nullptr);
 }
 
 bool UpdateEngineClientImpl::GetChannel(string* out_channel) {
-  return proxy_->GetChannel(
-      true,  // Get the current channel.
-      out_channel,
-      nullptr);
+  return proxy_->GetChannel(true,  // Get the current channel.
+                            out_channel, nullptr);
 }
 
 }  // namespace internal
diff --git a/client_library/client_impl.h b/client_library/client_impl.h
index 4886692..03c8f3c 100644
--- a/client_library/client_impl.h
+++ b/client_library/client_impl.h
@@ -67,9 +67,23 @@
 
   bool GetChannel(std::string* out_channel) override;
 
+  void RegisterStatusUpdateHandler(StatusUpdateHandler* handler) override;
+
  private:
   std::unique_ptr<org::chromium::UpdateEngineInterfaceProxy> proxy_;
 
+  void StatusUpdateHandlerRegistered(StatusUpdateHandler* handler,
+                                     const std::string& interface,
+                                     const std::string& signal_name,
+                                     bool success);
+
+  void RunStatusUpdateHandler(StatusUpdateHandler* handler,
+                              int64_t last_checked_time,
+                              double progress,
+                              const std::string& current_operation,
+                              const std::string& new_version,
+                              int64_t new_size);
+
   DISALLOW_COPY_AND_ASSIGN(UpdateEngineClientImpl);
 };  // class UpdateEngineClientImpl
 
diff --git a/client_library/include/update_engine/client.h b/client_library/include/update_engine/client.h
index 604cc5a..4d264c8 100644
--- a/client_library/include/update_engine/client.h
+++ b/client_library/include/update_engine/client.h
@@ -21,6 +21,7 @@
 #include <memory>
 #include <string>
 
+#include "update_engine/status_update_handler.h"
 #include "update_engine/update_status.h"
 
 namespace update_engine {
@@ -100,6 +101,13 @@
   // Get the channel the device is currently on.
   virtual bool GetChannel(std::string* out_channel) = 0;
 
+  // Handle status updates. The handler must exist until the client is
+  // destroyed. Its IPCError method will be called if the handler could
+  // not be registered. Otherwise its HandleStatusUpdate method will be called
+  // every time update_engine's status changes. Will always report the status
+  // on registration to prevent race conditions.
+  virtual void RegisterStatusUpdateHandler(StatusUpdateHandler* handler) = 0;
+
  protected:
   // Use CreateInstance().
   UpdateEngineClient() = default;
diff --git a/client_library/include/update_engine/status_update_handler.h b/client_library/include/update_engine/status_update_handler.h
new file mode 100644
index 0000000..d5b8cdb
--- /dev/null
+++ b/client_library/include/update_engine/status_update_handler.h
@@ -0,0 +1,47 @@
+//
+// Copyright (C) 2012 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.
+//
+
+#ifndef UPDATE_ENGINE_CLIENT_LIBRARY_INCLUDE_STATUS_UPDATE_HANDLER_H_
+#define UPDATE_ENGINE_CLIENT_LIBRARY_INCLUDE_STATUS_UPDATE_HANDLER_H_
+
+#include <string>
+
+#include "update_engine/client.h"
+#include "update_engine/update_status.h"
+
+namespace update_engine {
+
+// Handles update_engine status changes. An instance of this class can be
+// registered with UpdateEngineClient and will respond to any update_engine
+// status changes.
+class StatusUpdateHandler {
+ public:
+  virtual ~StatusUpdateHandler() = default;
+
+  // Runs when we fail to register the handler due to an IPC error.
+  virtual void IPCError(const std::string& error) = 0;
+
+  // Runs every time update_engine reports a status change.
+  virtual void HandleStatusUpdate(int64_t last_checked_time,
+                                  double progress,
+                                  UpdateStatus current_operation,
+                                  const std::string& new_version,
+                                  int64_t new_size) = 0;
+};
+
+}  // namespace update_engine
+
+#endif  // UPDATE_ENGINE_CLIENT_LIBRARY_INCLUDE_STATUS_UPDATE_HANDLER_H_
diff --git a/update_engine_client.cc b/update_engine_client.cc
index 9294d36..b5c9c24 100644
--- a/update_engine_client.cc
+++ b/update_engine_client.cc
@@ -20,6 +20,7 @@
 
 #include <memory>
 #include <string>
+#include <vector>
 
 #include <base/bind.h>
 #include <base/command_line.h>
@@ -30,15 +31,18 @@
 #include <dbus/bus.h>
 
 #include "update_engine/client.h"
-#include "update_engine/update_status.h"
 #include "update_engine/dbus-constants.h"
 #include "update_engine/dbus-proxies.h"
+#include "update_engine/status_update_handler.h"
+#include "update_engine/update_status.h"
 #include "update_status_utils.h"
 
 using std::string;
 using std::unique_ptr;
+using std::vector;
 using update_engine::kAttemptUpdateFlagNonInteractive;
 using update_engine::kUpdateEngineServiceName;
+using update_engine::UpdateStatus;
 using chromeos_update_engine::UpdateStatusToString;
 
 namespace {
@@ -58,18 +62,14 @@
  protected:
   int OnInit() override {
     int ret = DBusDaemon::OnInit();
-    if (ret != EX_OK)
-      return ret;
-    if (!InitProxy())
-      return 1;
+    if (ret != EX_OK) return ret;
+    if (!InitProxy()) return 1;
     // Wait for the UpdateEngine to be available or timeout.
-    proxy_->GetObjectProxy()->WaitForServiceToBeAvailable(
-        base::Bind(&UpdateEngineClient::OnServiceAvailable,
-                   base::Unretained(this)));
+    proxy_->GetObjectProxy()->WaitForServiceToBeAvailable(base::Bind(
+        &UpdateEngineClient::OnServiceAvailable, base::Unretained(this)));
     base::MessageLoop::current()->PostDelayedTask(
-        FROM_HERE,
-        base::Bind(&UpdateEngineClient::OnServiceAvailableTimeout,
-                   base::Unretained(this)),
+        FROM_HERE, base::Bind(&UpdateEngineClient::OnServiceAvailableTimeout,
+                              base::Unretained(this)),
         base::TimeDelta::FromSeconds(10));
     return EX_OK;
   }
@@ -84,23 +84,6 @@
   // after a timeout.
   void OnServiceAvailableTimeout();
 
-
-  // Callback called when a StatusUpdate signal is received.
-  void OnStatusUpdateSignal(int64_t last_checked_time,
-                            double progress,
-                            const string& current_operation,
-                            const string& new_version,
-                            int64_t new_size);
-  // Callback called when the OnStatusUpdateSignal() handler is registered.
-  void OnStatusUpdateSignalRegistration(const string& interface,
-                                        const string& signal_name,
-                                        bool success);
-
-  // Registers a callback that prints on stderr the received StatusUpdate
-  // signals.
-  // The daemon should continue running for this to work.
-  void WatchForUpdates();
-
   // Show the status of the update engine in stdout.
   bool ShowStatus();
 
@@ -108,32 +91,6 @@
   // occurred, 2 if no reboot is needed.
   int GetNeedReboot();
 
-  // This is similar to watching for updates but rather than registering
-  // a signal watch, actively poll the daemon just in case it stops
-  // sending notifications.
-  void WaitForUpdateComplete();
-  void OnUpdateCompleteCheck(int64_t last_checked_time,
-                             double progress,
-                             const string& current_operation,
-                             const string& new_version,
-                             int64_t new_size);
-
-  // Blocks until a reboot is needed. If the reboot is needed, exits the program
-  // with 0. Otherwise it exits the program with 1 if an error occurs before
-  // the reboot is needed.
-  void WaitForRebootNeeded();
-  void OnRebootNeededCheck(int64_t last_checked_time,
-                           double progress,
-                           const string& current_operation,
-                           const string& new_version,
-                           int64_t new_size);
-  // Callback called when the OnRebootNeededCheck() handler is registered. This
-  // is useful to check at this point if the reboot is needed, without loosing
-  // any StatusUpdate signals and avoiding the race condition.
-  void OnRebootNeededCheckRegistration(const string& interface,
-                                       const string& signal_name,
-                                       bool success);
-
   // Main method that parses and triggers all the actions based on the passed
   // flags.
   int ProcessFlags();
@@ -148,6 +105,9 @@
   // Library-based client
   unique_ptr<update_engine::UpdateEngineClient> client_;
 
+  // Pointers to handlers for cleanup
+  vector<unique_ptr<update_engine::StatusUpdateHandler>> handlers_;
+
   // Tell whether the UpdateEngine service is available after startup.
   bool service_is_available_{false};
 
@@ -171,8 +131,7 @@
     QuitWithExitCode(-1);
   }
   int ret = ProcessFlags();
-  if (ret != kContinueRunning)
-    QuitWithExitCode(ret);
+  if (ret != kContinueRunning) QuitWithExitCode(ret);
 }
 
 void UpdateEngineClient::OnServiceAvailableTimeout() {
@@ -183,59 +142,57 @@
   }
 }
 
-void UpdateEngineClient::OnStatusUpdateSignal(
-    int64_t last_checked_time,
-    double progress,
-    const string& current_operation,
-    const string& new_version,
-    int64_t new_size) {
+class ExitingStatusUpdateHandler : public update_engine::StatusUpdateHandler {
+ public:
+  ~ExitingStatusUpdateHandler() override = default;
+
+  void IPCError(const string& error) override;
+};
+
+void ExitingStatusUpdateHandler::IPCError(const string& error) {
+  LOG(ERROR) << error;
+  exit(1);
+}
+
+class WatchingStatusUpdateHandler : public ExitingStatusUpdateHandler {
+ public:
+  ~WatchingStatusUpdateHandler() override = default;
+
+  void HandleStatusUpdate(int64_t last_checked_time, double progress,
+                          UpdateStatus current_operation,
+                          const std::string& new_version,
+                          int64_t new_size) override;
+};
+
+void WatchingStatusUpdateHandler::HandleStatusUpdate(
+    int64_t last_checked_time, double progress, UpdateStatus current_operation,
+    const string& new_version, int64_t new_size) {
   LOG(INFO) << "Got status update:";
   LOG(INFO) << "  last_checked_time: " << last_checked_time;
   LOG(INFO) << "  progress: " << progress;
-  LOG(INFO) << "  current_operation: " << current_operation;
+  LOG(INFO) << "  current_operation: "
+            << UpdateStatusToString(current_operation);
   LOG(INFO) << "  new_version: " << new_version;
   LOG(INFO) << "  new_size: " << new_size;
 }
 
-void UpdateEngineClient::OnStatusUpdateSignalRegistration(
-    const string& interface,
-    const string& signal_name,
-    bool success) {
-  VLOG(1) << "OnStatusUpdateSignalRegistration(" << interface << ", "
-          << signal_name << ", " << success << ");";
-  if (!success) {
-    LOG(ERROR) << "Couldn't connect to the " << signal_name << " signal.";
-    exit(1);
-  }
-}
-
-void UpdateEngineClient::WatchForUpdates() {
-  proxy_->RegisterStatusUpdateSignalHandler(
-      base::Bind(&UpdateEngineClient::OnStatusUpdateSignal,
-                 base::Unretained(this)),
-      base::Bind(&UpdateEngineClient::OnStatusUpdateSignalRegistration,
-                 base::Unretained(this)));
-}
-
 bool UpdateEngineClient::ShowStatus() {
   int64_t last_checked_time = 0;
   double progress = 0.0;
-  update_engine::UpdateStatus current_op;
+  UpdateStatus current_op;
   string new_version;
   int64_t new_size = 0;
 
-  if (!client_->GetStatus(
-      &last_checked_time, &progress, &current_op, &new_version, &new_size)) {
+  if (!client_->GetStatus(&last_checked_time, &progress, &current_op,
+                          &new_version, &new_size)) {
     return false;
   }
 
-  printf("LAST_CHECKED_TIME=%" PRIi64 "\nPROGRESS=%f\nCURRENT_OP=%s\n"
+  printf("LAST_CHECKED_TIME=%" PRIi64
+         "\nPROGRESS=%f\nCURRENT_OP=%s\n"
          "NEW_VERSION=%s\nNEW_SIZE=%" PRIi64 "\n",
-         last_checked_time,
-         progress,
-         UpdateStatusToString(current_op),
-         new_version.c_str(),
-         new_size);
+         last_checked_time, progress, UpdateStatusToString(current_op),
+         new_version.c_str(), new_size);
 
   return true;
 }
@@ -243,98 +200,53 @@
 int UpdateEngineClient::GetNeedReboot() {
   int64_t last_checked_time = 0;
   double progress = 0.0;
-  update_engine::UpdateStatus current_op;
+  UpdateStatus current_op;
   string new_version;
   int64_t new_size = 0;
 
-  if (!client_->GetStatus(
-      &last_checked_time, &progress, &current_op, &new_version, &new_size)) {
+  if (!client_->GetStatus(&last_checked_time, &progress, &current_op,
+                          &new_version, &new_size)) {
     return 1;
   }
 
-  if (current_op == update_engine::UpdateStatus::UPDATED_NEED_REBOOT) {
+  if (current_op == UpdateStatus::UPDATED_NEED_REBOOT) {
     return 0;
   }
 
   return 2;
 }
 
-void UpdateEngineClient::OnUpdateCompleteCheck(
-    int64_t /* last_checked_time */,
-    double /* progress */,
-    const string& current_operation,
-    const string& /* new_version */,
-    int64_t /* new_size */) {
-  if (current_operation == update_engine::kUpdateStatusIdle) {
-    LOG(ERROR) << "Update failed, current operations is " << current_operation;
+class UpdateWaitHandler : public ExitingStatusUpdateHandler {
+ public:
+  UpdateWaitHandler(bool exit_on_error) : exit_on_error_(exit_on_error) {}
+
+  ~UpdateWaitHandler() override = default;
+
+  void HandleStatusUpdate(int64_t last_checked_time, double progress,
+                          UpdateStatus current_operation,
+                          const std::string& new_version,
+                          int64_t new_size) override;
+
+ private:
+  bool exit_on_error_;
+};
+
+void UpdateWaitHandler::HandleStatusUpdate(int64_t /* last_checked_time */,
+                                           double /* progress */,
+                                           UpdateStatus current_operation,
+                                           const string& /* new_version */,
+                                           int64_t /* new_size */) {
+  if (exit_on_error_ && current_operation == UpdateStatus::IDLE) {
+    LOG(ERROR) << "Update failed, current operations is "
+               << UpdateStatusToString(current_operation);
     exit(1);
   }
-  if (current_operation == update_engine::kUpdateStatusUpdatedNeedReboot) {
+  if (current_operation == UpdateStatus::UPDATED_NEED_REBOOT) {
     LOG(INFO) << "Update succeeded -- reboot needed.";
     exit(0);
   }
 }
 
-void UpdateEngineClient::WaitForUpdateComplete() {
-  proxy_->RegisterStatusUpdateSignalHandler(
-      base::Bind(&UpdateEngineClient::OnUpdateCompleteCheck,
-                 base::Unretained(this)),
-      base::Bind(&UpdateEngineClient::OnStatusUpdateSignalRegistration,
-                 base::Unretained(this)));
-}
-
-void UpdateEngineClient::OnRebootNeededCheck(
-    int64_t /* last_checked_time */,
-    double /* progress */,
-    const string& current_operation,
-    const string& /* new_version */,
-    int64_t /* new_size */) {
-  if (current_operation == update_engine::kUpdateStatusUpdatedNeedReboot) {
-    LOG(INFO) << "Reboot needed.";
-    exit(0);
-  }
-}
-
-void UpdateEngineClient::OnRebootNeededCheckRegistration(
-    const string& interface,
-    const string& signal_name,
-    bool success) {
-  int got = GetNeedReboot();
-
-  if (got == 1) {
-    LOG(ERROR) << "Could not query the current operation.";
-    exit(1);
-  }
-
-  if (got == 0) {
-    exit(0);
-  }
-
-  if (!success) {
-    LOG(ERROR) << "Couldn't connect to the " << signal_name << " signal.";
-    exit(1);
-  }
-}
-
-// Blocks until a reboot is needed. Returns true if waiting succeeded,
-// false if an error occurred.
-void UpdateEngineClient::WaitForRebootNeeded() {
-  proxy_->RegisterStatusUpdateSignalHandler(
-      base::Bind(&UpdateEngineClient::OnUpdateCompleteCheck,
-                 base::Unretained(this)),
-      base::Bind(&UpdateEngineClient::OnStatusUpdateSignalRegistration,
-                 base::Unretained(this)));
-  int got = GetNeedReboot();
-
-  if (got == 1) {
-    LOG(ERROR) << "Could not query the current operation.";
-    exit(1);
-  }
-
-  if (got == 0)
-    exit(0);
-}
-
 int UpdateEngineClient::ProcessFlags() {
   DEFINE_string(app_version, "", "Force the current app version.");
   DEFINE_string(channel, "",
@@ -342,25 +254,30 @@
                 "target channel is more stable than the current channel unless "
                 "--nopowerwash is specified.");
   DEFINE_bool(check_for_update, false, "Initiate check for updates.");
-  DEFINE_bool(follow, false, "Wait for any update operations to complete."
+  DEFINE_bool(follow, false,
+              "Wait for any update operations to complete."
               "Exit status is 0 if the update succeeded, and 1 otherwise.");
   DEFINE_bool(interactive, true, "Mark the update request as interactive.");
   DEFINE_string(omaha_url, "", "The URL of the Omaha update server.");
   DEFINE_string(p2p_update, "",
                 "Enables (\"yes\") or disables (\"no\") the peer-to-peer update"
                 " sharing.");
-  DEFINE_bool(powerwash, true, "When performing rollback or channel change, "
+  DEFINE_bool(powerwash, true,
+              "When performing rollback or channel change, "
               "do a powerwash or allow it respectively.");
   DEFINE_bool(reboot, false, "Initiate a reboot if needed.");
-  DEFINE_bool(is_reboot_needed, false, "Exit status 0 if reboot is needed, "
+  DEFINE_bool(is_reboot_needed, false,
+              "Exit status 0 if reboot is needed, "
               "2 if reboot is not needed or 1 if an error occurred.");
-  DEFINE_bool(block_until_reboot_is_needed, false, "Blocks until reboot is "
+  DEFINE_bool(block_until_reboot_is_needed, false,
+              "Blocks until reboot is "
               "needed. Returns non-zero exit status if an error occurred.");
   DEFINE_bool(reset_status, false, "Sets the status in update_engine to idle.");
   DEFINE_bool(rollback, false,
               "Perform a rollback to the previous partition. The device will "
               "be powerwashed unless --nopowerwash is specified.");
-  DEFINE_bool(can_rollback, false, "Shows whether rollback partition "
+  DEFINE_bool(can_rollback, false,
+              "Shows whether rollback partition "
               "is available.");
   DEFINE_bool(show_channel, false, "Show the current and target channels.");
   DEFINE_bool(show_p2p_update, false,
@@ -368,7 +285,8 @@
   DEFINE_bool(show_update_over_cellular, false,
               "Show the current setting for updates over cellular networks.");
   DEFINE_bool(status, false, "Print the status to stdout.");
-  DEFINE_bool(update, false, "Forces an update and waits for it to complete. "
+  DEFINE_bool(update, false,
+              "Forces an update and waits for it to complete. "
               "Implies --follow.");
   DEFINE_string(update_over_cellular, "",
                 "Enables (\"yes\") or disables (\"no\") the updates over "
@@ -414,7 +332,7 @@
       LOG(ERROR) << "Unknown option: \"" << FLAGS_update_over_cellular
                  << "\". Please specify \"yes\" or \"no\".";
     } else {
-      if(!client_->SetUpdateOverCellularPermission(allowed)) {
+      if (!client_->SetUpdateOverCellularPermission(allowed)) {
         LOG(ERROR) << "Error setting the update over cellular setting.";
         return 1;
       }
@@ -521,9 +439,9 @@
   }
 
   bool do_update_request = FLAGS_check_for_update | FLAGS_update |
-      !FLAGS_app_version.empty() | !FLAGS_omaha_url.empty();
-  if (FLAGS_update)
-    FLAGS_follow = true;
+                           !FLAGS_app_version.empty() |
+                           !FLAGS_omaha_url.empty();
+  if (FLAGS_update) FLAGS_follow = true;
 
   if (do_update_request && FLAGS_rollback) {
     LOG(ERROR) << "Incompatible flags specified with rollback."
@@ -556,9 +474,9 @@
   }
 
   // These final options are all mutually exclusive with one another.
-  if (FLAGS_follow + FLAGS_watch_for_updates + FLAGS_reboot +
-      FLAGS_status + FLAGS_is_reboot_needed +
-      FLAGS_block_until_reboot_is_needed > 1) {
+  if (FLAGS_follow + FLAGS_watch_for_updates + FLAGS_reboot + FLAGS_status +
+          FLAGS_is_reboot_needed + FLAGS_block_until_reboot_is_needed >
+      1) {
     LOG(ERROR) << "Multiple exclusive options selected. "
                << "Select only one of --follow, --watch_for_updates, --reboot, "
                << "--is_reboot_needed, --block_until_reboot_is_needed, "
@@ -577,13 +495,17 @@
 
   if (FLAGS_follow) {
     LOG(INFO) << "Waiting for update to complete.";
-    WaitForUpdateComplete();
+    auto handler = new UpdateWaitHandler(true);
+    handlers_.emplace_back(handler);
+    client_->RegisterStatusUpdateHandler(handler);
     return kContinueRunning;
   }
 
   if (FLAGS_watch_for_updates) {
     LOG(INFO) << "Watching for status updates.";
-    WatchForUpdates();
+    auto handler = new WatchingStatusUpdateHandler();
+    handlers_.emplace_back(handler);
+    client_->RegisterStatusUpdateHandler(handler);
     return kContinueRunning;
   }
 
@@ -614,7 +536,9 @@
   }
 
   if (FLAGS_block_until_reboot_is_needed) {
-    WaitForRebootNeeded();
+    auto handler = new UpdateWaitHandler(false);
+    handlers_.emplace_back(handler);
+    client_->RegisterStatusUpdateHandler(handler);
     return kContinueRunning;
   }