diff --git a/src/org/lineageos/updater/UpdaterReceiver.java b/src/org/lineageos/updater/UpdaterReceiver.java
index 6ef44056..42b60c64 100644
--- a/src/org/lineageos/updater/UpdaterReceiver.java
+++ b/src/org/lineageos/updater/UpdaterReceiver.java
@@ -91,7 +91,7 @@ public class UpdaterReceiver extends BroadcastReceiver {
             pm.reboot(null);
         } else if (Intent.ACTION_BOOT_COMPLETED.equals(intent.getAction())) {
             SharedPreferences pref = PreferenceManager.getDefaultSharedPreferences(context);
-            pref.edit().remove(Constants.PREF_NEEDS_REBOOT).apply();
+            pref.edit().remove(Constants.PREF_NEEDS_REBOOT_ID).apply();
 
             if (shouldShowUpdateFailedNotification(context)) {
                 pref.edit().putBoolean(Constants.PREF_INSTALL_NOTIFIED, true).apply();
diff --git a/src/org/lineageos/updater/UpdatesListAdapter.java b/src/org/lineageos/updater/UpdatesListAdapter.java
index 9d8a8545..e7814d6e 100644
--- a/src/org/lineageos/updater/UpdatesListAdapter.java
+++ b/src/org/lineageos/updater/UpdatesListAdapter.java
@@ -75,6 +75,7 @@ public class UpdatesListAdapter extends RecyclerView.Adapter<UpdatesListAdapter.
         INFO,
         DELETE,
         CANCEL_INSTALLATION,
+        REBOOT,
     }
 
     public static class ViewHolder extends RecyclerView.ViewHolder {
@@ -157,6 +158,11 @@ public class UpdatesListAdapter extends RecyclerView.Adapter<UpdatesListAdapter.
             setButtonAction(viewHolder.mAction, Action.INSTALL, downloadId, false);
             viewHolder.mProgressText.setText(R.string.list_verifying_update);
             viewHolder.mProgressBar.setIndeterminate(true);
+        } else if (mUpdaterController.isWaitingForReboot(downloadId)) {
+            setButtonAction(viewHolder.mAction, Action.REBOOT, downloadId, false);
+            viewHolder.mProgressText.setText(R.string.installing_update_finished);
+            viewHolder.mProgressBar.setIndeterminate(false);
+            viewHolder.mProgressBar.setProgress(100);
         } else {
             canDelete = true;
             setButtonAction(viewHolder.mAction, Action.RESUME, downloadId, !isBusy());
@@ -370,6 +376,13 @@ public class UpdatesListAdapter extends RecyclerView.Adapter<UpdatesListAdapter.
                 clickListener = enabled ? view -> getCancelInstallationDialog().show() : null;
             }
             break;
+            case REBOOT: {
+                button.setText(R.string.reboot);
+                button.setEnabled(enabled);
+                clickListener = enabled ?
+                        view -> mActivity.sendBroadcast(new Intent(Intent.ACTION_REBOOT)) : null;
+            }
+            break;
             default:
                 clickListener = null;
         }
diff --git a/src/org/lineageos/updater/controller/ABUpdateInstaller.java b/src/org/lineageos/updater/controller/ABUpdateInstaller.java
index 7012774c..82215f42 100644
--- a/src/org/lineageos/updater/controller/ABUpdateInstaller.java
+++ b/src/org/lineageos/updater/controller/ABUpdateInstaller.java
@@ -20,6 +20,7 @@ import android.content.SharedPreferences;
 import android.os.UpdateEngine;
 import android.os.UpdateEngineCallback;
 import android.support.v7.preference.PreferenceManager;
+import android.text.TextUtils;
 import android.util.Log;
 
 import org.lineageos.updater.misc.Constants;
@@ -117,13 +118,19 @@ class ABUpdateInstaller {
     static synchronized boolean isInstallingUpdate(Context context) {
         SharedPreferences pref = PreferenceManager.getDefaultSharedPreferences(context);
         return pref.getString(ABUpdateInstaller.PREF_INSTALLING_AB_ID, null) != null ||
-                pref.getBoolean(Constants.PREF_NEEDS_REBOOT, false);
+                pref.getString(Constants.PREF_NEEDS_REBOOT_ID, null) != null;
     }
 
     static synchronized boolean isInstallingUpdate(Context context, String downloadId) {
         SharedPreferences pref = PreferenceManager.getDefaultSharedPreferences(context);
         return downloadId.equals(pref.getString(ABUpdateInstaller.PREF_INSTALLING_AB_ID, null)) ||
-                pref.getBoolean(Constants.PREF_NEEDS_REBOOT, false);
+                TextUtils.equals(pref.getString(Constants.PREF_NEEDS_REBOOT_ID, null), downloadId);
+    }
+
+    static synchronized boolean isWaitingForReboot(Context context, String downloadId) {
+        String waitingId = PreferenceManager.getDefaultSharedPreferences(context)
+                .getString(Constants.PREF_NEEDS_REBOOT_ID, null);
+        return TextUtils.equals(waitingId, downloadId);
     }
 
     private ABUpdateInstaller(Context context, UpdaterController updaterController) {
@@ -230,8 +237,10 @@ class ABUpdateInstaller {
     }
 
     private void installationDone(boolean needsReboot) {
+        SharedPreferences prefs = PreferenceManager.getDefaultSharedPreferences(mContext);
+        String id = needsReboot ? prefs.getString(PREF_INSTALLING_AB_ID, null) : null;
         PreferenceManager.getDefaultSharedPreferences(mContext).edit()
-                .putBoolean(Constants.PREF_NEEDS_REBOOT, needsReboot)
+                .putString(Constants.PREF_NEEDS_REBOOT_ID, id)
                 .remove(PREF_INSTALLING_AB_ID)
                 .apply();
     }
diff --git a/src/org/lineageos/updater/misc/Constants.java b/src/org/lineageos/updater/misc/Constants.java
index 49106d4a..71582ae3 100644
--- a/src/org/lineageos/updater/misc/Constants.java
+++ b/src/org/lineageos/updater/misc/Constants.java
@@ -27,7 +27,7 @@ public final class Constants {
     public static final String PREF_AUTO_UPDATES_CHECK = "auto_updates_check";
     public static final String PREF_AUTO_DELETE_UPDATES = "auto_delete_updates";
     public static final String PREF_MOBILE_DATA_WARNING = "pref_mobile_data_warning";
-    public static final String PREF_NEEDS_REBOOT = "needs_reboot";
+    public static final String PREF_NEEDS_REBOOT_ID = "needs_reboot_id";
 
     public static final String UNCRYPT_FILE_EXT = ".uncrypt";