diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
index bf673c2..c4c8c9c 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
@@ -25,7 +25,6 @@
 import com.example.android.systemupdatersample.util.PayloadSpecs;
 import com.example.android.systemupdatersample.util.UpdateEngineErrorCodes;
 import com.example.android.systemupdatersample.util.UpdateEngineProperties;
-import com.example.android.systemupdatersample.util.UpdaterStates;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.AtomicDouble;
@@ -59,7 +58,7 @@
             new AtomicInteger(UpdateEngine.UpdateStatusConstants.IDLE);
     private AtomicInteger mEngineErrorCode = new AtomicInteger(UpdateEngineErrorCodes.UNKNOWN);
     private AtomicDouble mProgress = new AtomicDouble(0);
-    private AtomicInteger mState = new AtomicInteger(UpdaterStates.IDLE);
+    private UpdaterState mUpdaterState = new UpdaterState(UpdaterState.IDLE);
 
     private AtomicBoolean mManualSwitchSlotRequired = new AtomicBoolean(true);
 
@@ -111,7 +110,7 @@
 
     /**
      * Sets SystemUpdaterSample app state change callback. Value of {@code state} will be one
-     * of the values from {@link UpdaterStates}.
+     * of the values from {@link UpdaterState}.
      *
      * @param onStateChangeCallback a callback with parameter {@code state}.
      */
@@ -193,8 +192,14 @@
      * it also notifies {@link this.mOnStateChangeCallback}.
      */
     private void setUpdaterState(int updaterState) {
-        int previousState = mState.get();
-        mState.set(updaterState);
+        int previousState = mUpdaterState.get();
+        try {
+            mUpdaterState.set(updaterState);
+        } catch (UpdaterState.InvalidTransitionException e) {
+            // Note: invalid state transitions should be handled properly,
+            //       but to make sample app simple, we just throw runtime exception.
+            throw new RuntimeException("Can't set state " + updaterState, e);
+        }
         if (previousState != updaterState) {
             getOnStateChangeCallback().ifPresent(callback -> callback.accept(updaterState));
         }
@@ -211,7 +216,7 @@
     public void cancelRunningUpdate() {
         try {
             mUpdateEngine.cancel();
-            setUpdaterState(UpdaterStates.IDLE);
+            setUpdaterState(UpdaterState.IDLE);
         } catch (Exception e) {
             Log.w(TAG, "UpdateEngine failed to stop the ongoing update", e);
         }
@@ -227,7 +232,7 @@
     public void resetUpdate() {
         try {
             mUpdateEngine.resetStatus();
-            setUpdaterState(UpdaterStates.IDLE);
+            setUpdaterState(UpdaterState.IDLE);
         } catch (Exception e) {
             Log.w(TAG, "UpdateEngine failed to reset the update", e);
         }
@@ -241,7 +246,7 @@
      */
     public void applyUpdate(Context context, UpdateConfig config) {
         mEngineErrorCode.set(UpdateEngineErrorCodes.UNKNOWN);
-        setUpdaterState(UpdaterStates.RUNNING);
+        setUpdaterState(UpdaterState.RUNNING);
 
         synchronized (mLock) {
             // Cleaning up previous update data.
@@ -269,7 +274,7 @@
             builder.setPayload(mPayloadSpecs.forNonStreaming(config.getUpdatePackageFile()));
         } catch (IOException e) {
             Log.e(TAG, "Error creating payload spec", e);
-            setUpdaterState(UpdaterStates.ERROR);
+            setUpdaterState(UpdaterState.ERROR);
             return;
         }
         updateEngineApplyPayload(builder.build());
@@ -290,7 +295,7 @@
                 updateEngineApplyPayload(builder.build());
             } else {
                 Log.e(TAG, "PrepareStreamingService failed, result code is " + code);
-                setUpdaterState(UpdaterStates.ERROR);
+                setUpdaterState(UpdaterState.ERROR);
             }
         });
     }
@@ -332,7 +337,7 @@
                     properties.toArray(new String[0]));
         } catch (Exception e) {
             Log.e(TAG, "UpdateEngine failed to apply the update", e);
-            setUpdaterState(UpdaterStates.ERROR);
+            setUpdaterState(UpdaterState.ERROR);
         }
     }
 
@@ -396,9 +401,11 @@
         mEngineErrorCode.set(errorCode);
         if (errorCode == UpdateEngine.ErrorCodeConstants.SUCCESS
                 || errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
-            setUpdaterState(UpdaterStates.FINISHED);
+            setUpdaterState(isManualSwitchSlotRequired()
+                    ? UpdaterState.SLOT_SWITCH_REQUIRED
+                    : UpdaterState.REBOOT_REQUIRED);
         } else if (errorCode != UpdateEngineErrorCodes.USER_CANCELLED) {
-            setUpdaterState(UpdaterStates.ERROR);
+            setUpdaterState(UpdaterState.ERROR);
         }
 
         getOnEngineCompleteCallback()
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
new file mode 100644
index 0000000..36a9098
--- /dev/null
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
@@ -0,0 +1,103 @@
+/*
+ * 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;
+
+import android.util.SparseArray;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Controls updater state.
+ */
+public class UpdaterState {
+
+    public static final int IDLE = 0;
+    public static final int ERROR = 1;
+    public static final int RUNNING = 2;
+    public static final int PAUSED = 3;
+    public static final int SLOT_SWITCH_REQUIRED = 4;
+    public static final int REBOOT_REQUIRED = 5;
+
+    private static final SparseArray<String> STATE_MAP = new SparseArray<>();
+
+    static {
+        STATE_MAP.put(0, "IDLE");
+        STATE_MAP.put(1, "ERROR");
+        STATE_MAP.put(2, "RUNNING");
+        STATE_MAP.put(3, "PAUSED");
+        STATE_MAP.put(4, "SLOT_SWITCH_REQUIRED");
+        STATE_MAP.put(5, "REBOOT_REQUIRED");
+    }
+
+    /**
+     * Allowed state transitions. It's a map: key is a state, value is a set of states that
+     * are allowed to transition to from key.
+     */
+    private static final ImmutableMap<Integer, ImmutableSet<Integer>> TRANSITIONS =
+            ImmutableMap.of(
+                    IDLE, ImmutableSet.of(RUNNING),
+                    RUNNING, ImmutableSet.of(ERROR, PAUSED, REBOOT_REQUIRED, SLOT_SWITCH_REQUIRED),
+                    PAUSED, ImmutableSet.of(RUNNING),
+                    SLOT_SWITCH_REQUIRED, ImmutableSet.of(ERROR)
+            );
+
+    private AtomicInteger mState;
+
+    public UpdaterState(int state) {
+        this.mState = new AtomicInteger(state);
+    }
+
+    /**
+     * Returns updater state.
+     */
+    public int get() {
+        return mState.get();
+    }
+
+    /**
+     * Sets the updater state.
+     *
+     * @throws InvalidTransitionException if transition is not allowed.
+     */
+    public void set(int newState) throws InvalidTransitionException {
+        int oldState = mState.get();
+        if (!TRANSITIONS.get(oldState).contains(newState)) {
+            throw new InvalidTransitionException(
+                    "Can't transition from " + oldState + " to " + newState);
+        }
+        mState.set(newState);
+    }
+
+    /**
+     * Converts status code to status name.
+     */
+    public static String getStateText(int state) {
+        return STATE_MAP.get(state);
+    }
+
+    /**
+     * Defines invalid state transition exception.
+     */
+    public static class InvalidTransitionException extends Exception {
+        public InvalidTransitionException(String msg) {
+            super(msg);
+        }
+    }
+}
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 9983fe3..0b571cc 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -33,11 +33,11 @@
 import com.example.android.systemupdatersample.R;
 import com.example.android.systemupdatersample.UpdateConfig;
 import com.example.android.systemupdatersample.UpdateManager;
+import com.example.android.systemupdatersample.UpdaterState;
 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.UpdateEngineStatuses;
-import com.example.android.systemupdatersample.util.UpdaterStates;
 
 import java.util.List;
 
@@ -192,7 +192,7 @@
     /**
      * Invoked when SystemUpdaterSample app state changes.
      * Value of {@code state} will be one of the
-     * values from {@link UpdaterStates}.
+     * values from {@link UpdaterState}.
      */
     private void onUpdaterStateChange(int state) {
         Log.i(TAG, "onUpdaterStateChange invoked state=" + state);
@@ -233,8 +233,8 @@
         runOnUiThread(() -> {
             Log.i(TAG,
                     "Completed - errorCode="
-                    + UpdateEngineErrorCodes.getCodeName(errorCode) + "/" + errorCode
-                    + " " + completionState);
+                            + UpdateEngineErrorCodes.getCodeName(errorCode) + "/" + errorCode
+                            + " " + completionState);
             setUiEngineErrorCode(errorCode);
             if (errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
                 // if update was successfully applied.
@@ -323,7 +323,7 @@
      * @param state updater sample state
      */
     private void setUiUpdaterState(int state) {
-        String stateText = UpdaterStates.getStateText(state);
+        String stateText = UpdaterState.getStateText(state);
         mTextViewUpdaterState.setText(stateText + "/" + state);
     }
 
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/UpdaterStates.java b/updater_sample/src/com/example/android/systemupdatersample/util/UpdaterStates.java
deleted file mode 100644
index fc20a79..0000000
--- a/updater_sample/src/com/example/android/systemupdatersample/util/UpdaterStates.java
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * 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;
-
-import android.util.SparseArray;
-
-/**
- * SystemUpdaterSample app state.
- */
-public class UpdaterStates {
-
-    public static final int IDLE = 0;
-    public static final int ERROR = 1;
-    public static final int RUNNING = 2;
-    public static final int PAUSED = 3;
-    public static final int FINISHED = 4;
-
-    private static final SparseArray<String> STATE_MAP = new SparseArray<>();
-
-    static {
-        STATE_MAP.put(0, "IDLE");
-        STATE_MAP.put(1, "ERROR");
-        STATE_MAP.put(2, "RUNNING");
-        STATE_MAP.put(3, "PAUSED");
-        STATE_MAP.put(4, "FINISHED");
-    }
-
-    /**
-     * converts status code to status name
-     */
-    public static String getStateText(int state) {
-        return STATE_MAP.get(state);
-    }
-
-    private UpdaterStates() {}
-}