diff --git a/src/com/android/customization/module/ThemePickerInjector.java b/src/com/android/customization/module/ThemePickerInjector.java
index 263b44e..93c4d5f 100644
--- a/src/com/android/customization/module/ThemePickerInjector.java
+++ b/src/com/android/customization/module/ThemePickerInjector.java
@@ -35,6 +35,7 @@
 import com.android.customization.model.theme.ThemeManager;
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository;
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor;
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer;
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel;
 import com.android.systemui.shared.quickaffordance.data.content.KeyguardQuickAffordanceProviderClient;
 import com.android.systemui.shared.quickaffordance.data.content.KeyguardQuickAffordanceProviderClientImpl;
@@ -49,6 +50,9 @@
 import com.android.wallpaper.picker.ImagePreviewFragment;
 import com.android.wallpaper.picker.LivePreviewFragment;
 import com.android.wallpaper.picker.PreviewFragment;
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer;
+
+import java.util.Map;
 
 import kotlinx.coroutines.Dispatchers;
 
@@ -66,6 +70,7 @@
     private KeyguardQuickAffordanceProviderClient mKeyguardQuickAffordanceProviderClient;
     private FragmentFactory mFragmentFactory;
     private BaseFlags mFlags;
+    private KeyguardQuickAffordanceSnapshotRestorer mKeyguardQuickAffordanceSnapshotRestorer;
 
     @Override
     public CustomizationSections getCustomizationSections(Activity activity) {
@@ -149,7 +154,8 @@
                     getKeyguardQuickAffordancePickerProviderClient(context);
             mKeyguardQuickAffordancePickerInteractor = new KeyguardQuickAffordancePickerInteractor(
                     new KeyguardQuickAffordancePickerRepository(client, Dispatchers.getIO()),
-                    client);
+                    client,
+                    () -> getKeyguardQuickAffordanceSnapshotRestorer(context));
         }
         return mKeyguardQuickAffordancePickerInteractor;
     }
@@ -163,7 +169,8 @@
             mKeyguardQuickAffordancePickerViewModelFactory =
                     new KeyguardQuickAffordancePickerViewModel.Factory(
                             context,
-                            getKeyguardQuickAffordancePickerInteractor(context));
+                            getKeyguardQuickAffordancePickerInteractor(context),
+                            getUndoInteractor(context));
         }
         return mKeyguardQuickAffordancePickerViewModelFactory;
     }
@@ -176,8 +183,26 @@
         return mFragmentFactory;
     }
 
+    @Override
+    public BaseFlags getFlags() {
+        if (mFlags == null) {
+            mFlags = new BaseFlags() {};
+        }
+
+        return mFlags;
+    }
+
+    @Override
+    public Map<Integer, SnapshotRestorer> getSnapshotRestorers(Context context) {
+        final Map<Integer, SnapshotRestorer> restorers = super.getSnapshotRestorers(context);
+        restorers.put(
+                KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER,
+                getKeyguardQuickAffordanceSnapshotRestorer(context));
+        return restorers;
+    }
+
     /** Returns the {@link KeyguardQuickAffordanceProviderClient}. */
-    public KeyguardQuickAffordanceProviderClient getKeyguardQuickAffordancePickerProviderClient(
+    protected KeyguardQuickAffordanceProviderClient getKeyguardQuickAffordancePickerProviderClient(
             Context context) {
         if (mKeyguardQuickAffordanceProviderClient == null) {
             mKeyguardQuickAffordanceProviderClient =
@@ -187,12 +212,24 @@
         return mKeyguardQuickAffordanceProviderClient;
     }
 
-    @Override
-    public BaseFlags getFlags() {
-        if (mFlags == null) {
-            mFlags = new BaseFlags() {};
+    protected KeyguardQuickAffordanceSnapshotRestorer getKeyguardQuickAffordanceSnapshotRestorer(
+            Context context) {
+        if (mKeyguardQuickAffordanceSnapshotRestorer == null) {
+            mKeyguardQuickAffordanceSnapshotRestorer = new KeyguardQuickAffordanceSnapshotRestorer(
+                    getKeyguardQuickAffordancePickerInteractor(context),
+                    getKeyguardQuickAffordancePickerProviderClient(context));
         }
 
-        return mFlags;
+        return mKeyguardQuickAffordanceSnapshotRestorer;
     }
+
+    private static final int KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER =
+            WallpaperPicker2Injector.MIN_SNAPSHOT_RESTORER_KEY;
+
+    /**
+     * When this injector is overridden, this is the minimal value that should be used by restorers
+     * returns in {@link #getSnapshotRestorers(Context)}.
+     */
+    protected static final int MIN_SNAPSHOT_RESTORER_KEY =
+            KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER + 1;
 }
diff --git a/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt b/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt
new file mode 100644
index 0000000..63b4a9b
--- /dev/null
+++ b/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2022 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.android.customization.picker.clock.domain.interactor
+
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+
+/** Handles state restoration for clocks. */
+class ClocksSnapshotRestorer : SnapshotRestorer {
+    override suspend fun setUpSnapshotRestorer(
+        updater: (RestorableSnapshot) -> Unit,
+    ): RestorableSnapshot {
+        // TODO(b/262924055): implement as part of the clock settings screen.
+        return RestorableSnapshot(mapOf())
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        // TODO(b/262924055): implement as part of the clock settings screen.
+    }
+}
diff --git a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt
index 87cedf5..c2d2e5a 100644
--- a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt
+++ b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt
@@ -24,6 +24,7 @@
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel as SelectionModel
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSlotModel as SlotModel
 import com.android.systemui.shared.quickaffordance.data.content.KeyguardQuickAffordanceProviderClient as Client
+import javax.inject.Provider
 import kotlinx.coroutines.flow.Flow
 
 /**
@@ -33,10 +34,8 @@
 class KeyguardQuickAffordancePickerInteractor(
     private val repository: KeyguardQuickAffordancePickerRepository,
     private val client: Client,
+    private val snapshotRestorer: Provider<KeyguardQuickAffordanceSnapshotRestorer>,
 ) {
-    /** Whether the feature is enabled. */
-    val isFeatureEnabled: Flow<Boolean> = repository.isFeatureEnabled
-
     /** List of slots available on the device. */
     val slots: Flow<List<SlotModel>> = repository.slots
 
@@ -60,6 +59,8 @@
             slotId = slotId,
             affordanceId = affordanceId,
         )
+
+        snapshotRestorer.get().storeSnapshot()
     }
 
     /** Unselects an affordance with the given ID from the slot with the given ID. */
@@ -68,6 +69,8 @@
             slotId = slotId,
             affordanceId = affordanceId,
         )
+
+        snapshotRestorer.get().storeSnapshot()
     }
 
     /** Unselects all affordances from the slot with the given ID. */
@@ -75,6 +78,8 @@
         client.deleteAllSelections(
             slotId = slotId,
         )
+
+        snapshotRestorer.get().storeSnapshot()
     }
 
     /** Returns a [Drawable] for the given resource ID, from the system UI package. */
diff --git a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
new file mode 100644
index 0000000..f0544a1
--- /dev/null
+++ b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2022 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.android.customization.picker.quickaffordance.domain.interactor
+
+import com.android.systemui.shared.quickaffordance.data.content.KeyguardQuickAffordanceProviderClient
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+
+/** Handles state restoration for the quick affordances system. */
+class KeyguardQuickAffordanceSnapshotRestorer(
+    private val interactor: KeyguardQuickAffordancePickerInteractor,
+    private val client: KeyguardQuickAffordanceProviderClient,
+) : SnapshotRestorer {
+
+    private lateinit var snapshotUpdater: (RestorableSnapshot) -> Unit
+
+    suspend fun storeSnapshot() {
+        snapshotUpdater(snapshot())
+    }
+
+    override suspend fun setUpSnapshotRestorer(
+        updater: (RestorableSnapshot) -> Unit,
+    ): RestorableSnapshot {
+        snapshotUpdater = updater
+        return snapshot()
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        val selections: List<Pair<String, String>> =
+            checkNotNull(snapshot.args[KEY_SELECTIONS]).split(SELECTION_SEPARATOR).map { selection
+                ->
+                val (slotId, affordanceId) = selection.split(SLOT_AFFORDANCE_SEPARATOR)
+                slotId to affordanceId
+            }
+
+        selections.forEach { (slotId, affordanceId) ->
+            interactor.select(
+                slotId,
+                affordanceId,
+            )
+        }
+    }
+
+    private suspend fun snapshot(): RestorableSnapshot {
+        return RestorableSnapshot(
+            mapOf(
+                KEY_SELECTIONS to
+                    client.querySelections().joinToString(SELECTION_SEPARATOR) { selection ->
+                        "${selection.slotId}${SLOT_AFFORDANCE_SEPARATOR}${selection.affordanceId}"
+                    }
+            )
+        )
+    }
+
+    companion object {
+        private const val KEY_SELECTIONS = "selections"
+        private const val SLOT_AFFORDANCE_SEPARATOR = "->"
+        private const val SELECTION_SEPARATOR = "|"
+    }
+}
diff --git a/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt b/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt
index 251cdca..f37b246 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt
@@ -30,6 +30,7 @@
 import com.android.wallpaper.R
 import com.android.wallpaper.module.InjectorProvider
 import com.android.wallpaper.picker.AppbarFragment
+import com.android.wallpaper.picker.undo.ui.binder.RevertToolbarButtonBinder
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.suspendCancellableCoroutine
 
@@ -63,6 +64,13 @@
                     injector.getKeyguardQuickAffordancePickerViewModelFactory(requireContext()),
                 )
                 .get()
+        setUpToolbarMenu(R.menu.undoable_customization_menu)
+        RevertToolbarButtonBinder.bind(
+            view = view.requireViewById(toolbarId),
+            viewModel = viewModel.undo,
+            lifecycleOwner = this,
+        )
+
         KeyguardQuickAffordancePreviewBinder.bind(
             activity = requireActivity(),
             previewView = view.requireViewById(R.id.preview),
diff --git a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
index 774ff22..fa43edd 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
@@ -29,6 +29,8 @@
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
 import com.android.systemui.shared.quickaffordance.data.content.KeyguardQuickAffordanceProviderContract as Contract
 import com.android.wallpaper.R
+import com.android.wallpaper.picker.undo.domain.interactor.UndoInteractor
+import com.android.wallpaper.picker.undo.ui.viewmodel.UndoViewModel
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -43,20 +45,26 @@
 class KeyguardQuickAffordancePickerViewModel
 private constructor(
     context: Context,
-    private val interactor: KeyguardQuickAffordancePickerInteractor,
+    private val quickAffordanceInteractor: KeyguardQuickAffordancePickerInteractor,
+    undoInteractor: UndoInteractor,
 ) : ViewModel() {
 
     @SuppressLint("StaticFieldLeak") private val applicationContext = context.applicationContext
 
+    val undo: UndoViewModel =
+        UndoViewModel(
+            interactor = undoInteractor,
+        )
+
     private val _selectedSlotId = MutableStateFlow<String?>(null)
     val selectedSlotId: StateFlow<String?> = _selectedSlotId.asStateFlow()
 
     /** View-models for each slot, keyed by slot ID. */
     val slots: Flow<Map<String, KeyguardQuickAffordanceSlotViewModel>> =
         combine(
-            interactor.slots,
-            interactor.affordances,
-            interactor.selections,
+            quickAffordanceInteractor.slots,
+            quickAffordanceInteractor.affordances,
+            quickAffordanceInteractor.selections,
             selectedSlotId,
         ) { slots, affordances, selections, selectedSlotIdOrNull ->
             slots
@@ -102,9 +110,9 @@
     /** The list of all available quick affordances for the selected slot. */
     val quickAffordances: Flow<List<KeyguardQuickAffordanceViewModel>> =
         combine(
-            interactor.slots,
-            interactor.affordances,
-            interactor.selections,
+            quickAffordanceInteractor.slots,
+            quickAffordanceInteractor.affordances,
+            quickAffordanceInteractor.selections,
             selectedSlotId,
         ) { slots, affordances, selections, selectedSlotIdOrNull ->
             val selectedSlot =
@@ -131,12 +139,12 @@
                                 {
                                     viewModelScope.launch {
                                         if (isSelected) {
-                                            interactor.unselect(
+                                            quickAffordanceInteractor.unselect(
                                                 slotId = selectedSlot.id,
                                                 affordanceId = affordance.id,
                                             )
                                         } else {
-                                            interactor.select(
+                                            quickAffordanceInteractor.select(
                                                 slotId = selectedSlot.id,
                                                 affordanceId = affordance.id,
                                             )
@@ -229,7 +237,9 @@
         return KeyguardQuickAffordanceViewModel.none(
             context = applicationContext,
             isSelected = isSelected,
-            onSelected = { viewModelScope.launch { interactor.unselectAll(slotId) } },
+            onSelected = {
+                viewModelScope.launch { quickAffordanceInteractor.unselectAll(slotId) }
+            },
         )
     }
 
@@ -246,7 +256,7 @@
     }
 
     private suspend fun getAffordanceIcon(@DrawableRes iconResourceId: Int): Drawable {
-        return interactor.getAffordanceIcon(iconResourceId)
+        return quickAffordanceInteractor.getAffordanceIcon(iconResourceId)
     }
 
     private fun String?.toIntent(): Intent? {
@@ -318,13 +328,15 @@
 
     class Factory(
         private val context: Context,
-        private val interactor: KeyguardQuickAffordancePickerInteractor,
+        private val quickAffordanceInteractor: KeyguardQuickAffordancePickerInteractor,
+        private val undoInteractor: UndoInteractor,
     ) : ViewModelProvider.Factory {
         override fun <T : ViewModel> create(modelClass: Class<T>): T {
             @Suppress("UNCHECKED_CAST")
             return KeyguardQuickAffordancePickerViewModel(
                 context = context,
-                interactor = interactor,
+                quickAffordanceInteractor = quickAffordanceInteractor,
+                undoInteractor = undoInteractor,
             )
                 as T
         }
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
index d8a136d..87f47fa 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
+++ b/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -20,16 +20,17 @@
 import androidx.test.filters.SmallTest
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
 import com.android.systemui.shared.quickaffordance.data.content.FakeKeyguardQuickAffordanceProviderClient
+import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.toList
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.resetMain
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.test.setMain
@@ -51,18 +52,25 @@
 
     @Before
     fun setUp() {
-        val coroutineDispatcher = UnconfinedTestDispatcher()
-        testScope = TestScope(coroutineDispatcher)
-        Dispatchers.setMain(coroutineDispatcher)
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        Dispatchers.setMain(testDispatcher)
         client = FakeKeyguardQuickAffordanceProviderClient()
         underTest =
             KeyguardQuickAffordancePickerInteractor(
                 repository =
                     KeyguardQuickAffordancePickerRepository(
                         client = client,
-                        backgroundDispatcher = coroutineDispatcher,
+                        backgroundDispatcher = testDispatcher,
                     ),
                 client = client,
+                snapshotRestorer = {
+                    KeyguardQuickAffordanceSnapshotRestorer(
+                            interactor = underTest,
+                            client = client,
+                        )
+                        .apply { runBlocking { setUpSnapshotRestorer {} } }
+                },
             )
     }
 
@@ -74,14 +82,13 @@
     @Test
     fun select() =
         testScope.runTest {
-            val selections = mutableListOf<List<KeyguardQuickAffordancePickerSelectionModel>>()
-            val job = launch { underTest.selections.toList(selections) }
+            val selections = collectLastValue(underTest.selections)
 
             underTest.select(
                 slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
                 affordanceId = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_1,
             )
-            assertThat(selections.last())
+            assertThat(selections())
                 .isEqualTo(
                     listOf(
                         KeyguardQuickAffordancePickerSelectionModel(
@@ -95,7 +102,7 @@
                 slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
                 affordanceId = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_2,
             )
-            assertThat(selections.last())
+            assertThat(selections())
                 .isEqualTo(
                     listOf(
                         KeyguardQuickAffordancePickerSelectionModel(
@@ -104,15 +111,12 @@
                         ),
                     )
                 )
-
-            job.cancel()
         }
 
     @Test
     fun unselect() =
         testScope.runTest {
-            val selections = mutableListOf<List<KeyguardQuickAffordancePickerSelectionModel>>()
-            val job = launch { underTest.selections.toList(selections) }
+            val selections = collectLastValue(underTest.selections)
             underTest.select(
                 slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
                 affordanceId = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_1,
@@ -123,17 +127,14 @@
                 affordanceId = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_1,
             )
 
-            assertThat(selections.last()).isEmpty()
-
-            job.cancel()
+            assertThat(selections()).isEmpty()
         }
 
     @Test
     fun unselectAll() =
         testScope.runTest {
             client.setSlotCapacity(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END, 3)
-            val selections = mutableListOf<List<KeyguardQuickAffordancePickerSelectionModel>>()
-            val job = launch { underTest.selections.toList(selections) }
+            val selections = collectLastValue(underTest.selections)
             underTest.select(
                 slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
                 affordanceId = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_1,
@@ -151,8 +152,6 @@
                 slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
             )
 
-            assertThat(selections.last()).isEmpty()
-
-            job.cancel()
+            assertThat(selections()).isEmpty()
         }
 }
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
index d446e1b..3ec893a 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
+++ b/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -22,6 +22,7 @@
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceSlotViewModel
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceSummaryViewModel
@@ -29,14 +30,17 @@
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
 import com.android.systemui.shared.quickaffordance.data.content.FakeKeyguardQuickAffordanceProviderClient
 import com.android.systemui.shared.quickaffordance.data.content.KeyguardQuickAffordanceProviderClient
+import com.android.wallpaper.picker.undo.data.repository.UndoRepository
+import com.android.wallpaper.picker.undo.domain.interactor.UndoInteractor
+import com.android.wallpaper.testing.FAKE_RESTORERS
+import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.toList
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.resetMain
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.test.setMain
@@ -56,27 +60,43 @@
     private lateinit var context: Context
     private lateinit var testScope: TestScope
     private lateinit var client: FakeKeyguardQuickAffordanceProviderClient
+    private lateinit var quickAffordanceInteractor: KeyguardQuickAffordancePickerInteractor
 
     @Before
     fun setUp() {
         context = InstrumentationRegistry.getInstrumentation().targetContext
-        val coroutineDispatcher = UnconfinedTestDispatcher()
-        testScope = TestScope(coroutineDispatcher)
-        Dispatchers.setMain(coroutineDispatcher)
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        Dispatchers.setMain(testDispatcher)
         client = FakeKeyguardQuickAffordanceProviderClient()
 
+        quickAffordanceInteractor =
+            KeyguardQuickAffordancePickerInteractor(
+                repository =
+                    KeyguardQuickAffordancePickerRepository(
+                        client = client,
+                        backgroundDispatcher = testDispatcher,
+                    ),
+                client = client,
+                snapshotRestorer = {
+                    KeyguardQuickAffordanceSnapshotRestorer(
+                            interactor = quickAffordanceInteractor,
+                            client = client,
+                        )
+                        .apply { runBlocking { setUpSnapshotRestorer {} } }
+                },
+            )
+        val undoInteractor =
+            UndoInteractor(
+                scope = testScope.backgroundScope,
+                repository = UndoRepository(),
+                restorerByOwnerId = FAKE_RESTORERS,
+            )
         underTest =
             KeyguardQuickAffordancePickerViewModel.Factory(
                     context = context,
-                    interactor =
-                        KeyguardQuickAffordancePickerInteractor(
-                            repository =
-                                KeyguardQuickAffordancePickerRepository(
-                                    client = client,
-                                    backgroundDispatcher = coroutineDispatcher,
-                                ),
-                            client = client,
-                        ),
+                    quickAffordanceInteractor = quickAffordanceInteractor,
+                    undoInteractor = undoInteractor,
                 )
                 .create(KeyguardQuickAffordancePickerViewModel::class.java)
     }
@@ -89,23 +109,18 @@
     @Test
     fun `Select an affordance for each side`() =
         testScope.runTest {
-            val slots = mutableListOf<Map<String, KeyguardQuickAffordanceSlotViewModel>>()
-            val quickAffordances = mutableListOf<List<KeyguardQuickAffordanceViewModel>>()
-
-            val jobs = buildList {
-                add(launch { underTest.slots.toList(slots) })
-                add(launch { underTest.quickAffordances.toList(quickAffordances) })
-            }
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
 
             // Initially, the first slot is selected with the "none" affordance selected.
             assertPickerUiState(
-                slots = slots.last(),
-                affordances = quickAffordances.last(),
+                slots = slots(),
+                affordances = quickAffordances(),
                 selectedSlotText = "Left button",
                 selectedAffordanceText = "None",
             )
             assertPreviewUiState(
-                slots = slots.last(),
+                slots = slots(),
                 expectedAffordanceNameBySlotId =
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to null,
@@ -114,15 +129,15 @@
             )
 
             // Select "affordance 1" for the first slot.
-            quickAffordances.last()[1].onClicked?.invoke()
+            quickAffordances()?.get(1)?.onClicked?.invoke()
             assertPickerUiState(
-                slots = slots.last(),
-                affordances = quickAffordances.last(),
+                slots = slots(),
+                affordances = quickAffordances(),
                 selectedSlotText = "Left button",
                 selectedAffordanceText = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_1,
             )
             assertPreviewUiState(
-                slots = slots.last(),
+                slots = slots(),
                 expectedAffordanceNameBySlotId =
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
@@ -133,17 +148,17 @@
 
             // Select an affordance for the second slot.
             // First, switch to the second slot:
-            slots.last()[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END]?.onClicked?.invoke()
+            slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances.last()[3].onClicked?.invoke()
+            quickAffordances()?.get(3)?.onClicked?.invoke()
             assertPickerUiState(
-                slots = slots.last(),
-                affordances = quickAffordances.last(),
+                slots = slots(),
+                affordances = quickAffordances(),
                 selectedSlotText = "Right button",
                 selectedAffordanceText = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_3,
             )
             assertPreviewUiState(
-                slots = slots.last(),
+                slots = slots(),
                 expectedAffordanceNameBySlotId =
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
@@ -154,15 +169,15 @@
             )
 
             // Select a different affordance for the second slot.
-            quickAffordances.last()[2].onClicked?.invoke()
+            quickAffordances()?.get(2)?.onClicked?.invoke()
             assertPickerUiState(
-                slots = slots.last(),
-                affordances = quickAffordances.last(),
+                slots = slots(),
+                affordances = quickAffordances(),
                 selectedSlotText = "Right button",
                 selectedAffordanceText = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_2,
             )
             assertPreviewUiState(
-                slots = slots.last(),
+                slots = slots(),
                 expectedAffordanceNameBySlotId =
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
@@ -171,42 +186,35 @@
                             FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_2,
                     ),
             )
-
-            jobs.forEach { it.cancel() }
         }
 
     @Test
     fun `Unselect - AKA selecting the none affordance - on one side`() =
         testScope.runTest {
-            val slots = mutableListOf<Map<String, KeyguardQuickAffordanceSlotViewModel>>()
-            val quickAffordances = mutableListOf<List<KeyguardQuickAffordanceViewModel>>()
-
-            val jobs = buildList {
-                add(launch { underTest.slots.toList(slots) })
-                add(launch { underTest.quickAffordances.toList(quickAffordances) })
-            }
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
 
             // Select "affordance 1" for the first slot.
-            quickAffordances.last()[1].onClicked?.invoke()
+            quickAffordances()?.get(1)?.onClicked?.invoke()
             // Select an affordance for the second slot.
             // First, switch to the second slot:
-            slots.last()[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END]?.onClicked?.invoke()
+            slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances.last()[3].onClicked?.invoke()
+            quickAffordances()?.get(3)?.onClicked?.invoke()
 
             // Switch back to the first slot:
-            slots.last()[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START]?.onClicked?.invoke()
+            slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START)?.onClicked?.invoke()
             // Select the "none" affordance, which is always in position 0:
-            quickAffordances.last()[0].onClicked?.invoke()
+            quickAffordances()?.get(0)?.onClicked?.invoke()
 
             assertPickerUiState(
-                slots = slots.last(),
-                affordances = quickAffordances.last(),
+                slots = slots(),
+                affordances = quickAffordances(),
                 selectedSlotText = "Left button",
                 selectedAffordanceText = "None",
             )
             assertPreviewUiState(
-                slots = slots.last(),
+                slots = slots(),
                 expectedAffordanceNameBySlotId =
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to null,
@@ -214,22 +222,15 @@
                             FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_3,
                     ),
             )
-
-            jobs.forEach { it.cancel() }
         }
 
     @Test
     fun `Show enablement dialog when selecting a disabled affordance`() =
         testScope.runTest {
-            val slots = mutableListOf<Map<String, KeyguardQuickAffordanceSlotViewModel>>()
-            val quickAffordances = mutableListOf<List<KeyguardQuickAffordanceViewModel>>()
-            val dialog = mutableListOf<KeyguardQuickAffordancePickerViewModel.DialogViewModel?>()
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val dialog = collectLastValue(underTest.dialog)
 
-            val jobs = buildList {
-                add(launch { underTest.slots.toList(slots) })
-                add(launch { underTest.quickAffordances.toList(quickAffordances) })
-                add(launch { underTest.dialog.toList(dialog) })
-            }
             val enablementInstructions = listOf("header", "enablementInstructions")
             val enablementActionText = "enablementActionText"
             val packageName = "packageName"
@@ -250,46 +251,39 @@
                 )
 
             // Lets try to select that disabled affordance:
-            quickAffordances.last()[affordanceIndex + 1].onClicked?.invoke()
+            quickAffordances()?.get(affordanceIndex + 1)?.onClicked?.invoke()
 
             // We expect there to be a dialog that should be shown:
-            assertThat(dialog.last()?.instructionHeader).isEqualTo(enablementInstructions[0])
-            assertThat(dialog.last()?.instructions)
+            assertThat(dialog()?.instructionHeader).isEqualTo(enablementInstructions[0])
+            assertThat(dialog()?.instructions)
                 .isEqualTo(enablementInstructions.subList(1, enablementInstructions.size))
-            assertThat(dialog.last()?.actionText).isEqualTo(enablementActionText)
-            assertThat(dialog.last()?.intent?.`package`).isEqualTo(packageName)
-            assertThat(dialog.last()?.intent?.action).isEqualTo(action)
+            assertThat(dialog()?.actionText).isEqualTo(enablementActionText)
+            assertThat(dialog()?.intent?.`package`).isEqualTo(packageName)
+            assertThat(dialog()?.intent?.action).isEqualTo(action)
 
             // Once we report that the dialog has been dismissed by the user, we expect there to be
             // no
             // dialog to be shown:
             underTest.onDialogDismissed()
-            assertThat(dialog.last()).isNull()
-
-            jobs.forEach { it.cancel() }
+            assertThat(dialog()).isNull()
         }
 
     @Test
     fun `summary - affordance selected in both bottom-start and bottom-end`() =
         testScope.runTest {
-            val slots = mutableListOf<Map<String, KeyguardQuickAffordanceSlotViewModel>>()
-            val quickAffordances = mutableListOf<List<KeyguardQuickAffordanceViewModel>>()
-            val summary = mutableListOf<KeyguardQuickAffordanceSummaryViewModel>()
-            val jobs = buildList {
-                add(launch { underTest.slots.toList(slots) })
-                add(launch { underTest.quickAffordances.toList(quickAffordances) })
-                add(launch { underTest.summary.toList(summary) })
-            }
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val summary = collectLastValue(underTest.summary)
 
             // Select "affordance 1" for the first slot.
-            quickAffordances.last()[1].onClicked?.invoke()
+            quickAffordances()?.get(1)?.onClicked?.invoke()
             // Select an affordance for the second slot.
             // First, switch to the second slot:
-            slots.last()[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END]?.onClicked?.invoke()
+            slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances.last()[3].onClicked?.invoke()
+            quickAffordances()?.get(3)?.onClicked?.invoke()
 
-            assertThat(summary.last())
+            assertThat(summary())
                 .isEqualTo(
                     KeyguardQuickAffordanceSummaryViewModel(
                         description =
@@ -299,25 +293,19 @@
                         icon2 = FakeKeyguardQuickAffordanceProviderClient.ICON_3,
                     )
                 )
-            jobs.forEach { it.cancel() }
         }
 
     @Test
     fun `summary - affordance selected only on bottom-start`() =
         testScope.runTest {
-            val slots = mutableListOf<Map<String, KeyguardQuickAffordanceSlotViewModel>>()
-            val quickAffordances = mutableListOf<List<KeyguardQuickAffordanceViewModel>>()
-            val summary = mutableListOf<KeyguardQuickAffordanceSummaryViewModel>()
-            val jobs = buildList {
-                add(launch { underTest.slots.toList(slots) })
-                add(launch { underTest.quickAffordances.toList(quickAffordances) })
-                add(launch { underTest.summary.toList(summary) })
-            }
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val summary = collectLastValue(underTest.summary)
 
             // Select "affordance 1" for the first slot.
-            quickAffordances.last()[1].onClicked?.invoke()
+            quickAffordances()?.get(1)?.onClicked?.invoke()
 
-            assertThat(summary.last())
+            assertThat(summary())
                 .isEqualTo(
                     KeyguardQuickAffordanceSummaryViewModel(
                         description = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_1,
@@ -325,28 +313,22 @@
                         icon2 = null,
                     )
                 )
-            jobs.forEach { it.cancel() }
         }
 
     @Test
     fun `summary - affordance selected only on bottom-end`() =
         testScope.runTest {
-            val slots = mutableListOf<Map<String, KeyguardQuickAffordanceSlotViewModel>>()
-            val quickAffordances = mutableListOf<List<KeyguardQuickAffordanceViewModel>>()
-            val summary = mutableListOf<KeyguardQuickAffordanceSummaryViewModel>()
-            val jobs = buildList {
-                add(launch { underTest.slots.toList(slots) })
-                add(launch { underTest.quickAffordances.toList(quickAffordances) })
-                add(launch { underTest.summary.toList(summary) })
-            }
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val summary = collectLastValue(underTest.summary)
 
             // Select an affordance for the second slot.
             // First, switch to the second slot:
-            slots.last()[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END]?.onClicked?.invoke()
+            slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances.last()[3].onClicked?.invoke()
+            quickAffordances()?.get(3)?.onClicked?.invoke()
 
-            assertThat(summary.last())
+            assertThat(summary())
                 .isEqualTo(
                     KeyguardQuickAffordanceSummaryViewModel(
                         description = FakeKeyguardQuickAffordanceProviderClient.AFFORDANCE_3,
@@ -354,25 +336,18 @@
                         icon2 = FakeKeyguardQuickAffordanceProviderClient.ICON_3,
                     )
                 )
-            jobs.forEach { it.cancel() }
         }
 
     @Test
     fun `summary - no affordances selected`() =
         testScope.runTest {
-            val slots = mutableListOf<Map<String, KeyguardQuickAffordanceSlotViewModel>>()
-            val quickAffordances = mutableListOf<List<KeyguardQuickAffordanceViewModel>>()
-            val summary = mutableListOf<KeyguardQuickAffordanceSummaryViewModel>()
-            val jobs = buildList {
-                add(launch { underTest.slots.toList(slots) })
-                add(launch { underTest.quickAffordances.toList(quickAffordances) })
-                add(launch { underTest.summary.toList(summary) })
-            }
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val summary = collectLastValue(underTest.summary)
 
-            assertThat(summary.last().description).isEqualTo("None")
-            assertThat(summary.last().icon1).isNotNull()
-            assertThat(summary.last().icon2).isNull()
-            jobs.forEach { it.cancel() }
+            assertThat(summary()?.description).isEqualTo("None")
+            assertThat(summary()?.icon1).isNotNull()
+            assertThat(summary()?.icon2).isNull()
         }
 
     /**
@@ -385,8 +360,8 @@
      * @param selectedAffordanceText The text of the affordance that's expected to be selected
      */
     private fun assertPickerUiState(
-        slots: Map<String, KeyguardQuickAffordanceSlotViewModel>,
-        affordances: List<KeyguardQuickAffordanceViewModel>,
+        slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
+        affordances: List<KeyguardQuickAffordanceViewModel>?,
         selectedSlotText: String,
         selectedAffordanceText: String,
     ) {
@@ -402,7 +377,8 @@
         )
 
         var foundSelectedAffordance = false
-        affordances.forEach { affordance ->
+        assertThat(affordances).isNotNull()
+        affordances?.forEach { affordance ->
             val nameMatchesSelectedName = affordance.contentDescription == selectedAffordanceText
             assertWithMessage(
                     "Expected affordance with name \"${affordance.contentDescription}\" to have" +
@@ -423,11 +399,11 @@
      * @param isSelected Whether that slot should be selected
      */
     private fun assertSlotTabUiState(
-        slots: Map<String, KeyguardQuickAffordanceSlotViewModel>,
+        slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
         slotId: String,
         isSelected: Boolean,
     ) {
-        val viewModel = slots[slotId] ?: error("No slot with ID \"$slotId\"!")
+        val viewModel = slots?.get(slotId) ?: error("No slot with ID \"$slotId\"!")
         assertThat(viewModel.isSelected).isEqualTo(isSelected)
     }
 
@@ -436,13 +412,15 @@
      *
      * @param slots The observed slot view-models, keyed by slot ID
      * @param expectedAffordanceNameBySlotId The expected name of the selected affordance for each
-     * slot ID or `null` if it's expected for there to be no affordance for that slot in the preview
+     *   slot ID or `null` if it's expected for there to be no affordance for that slot in the
+     *   preview
      */
     private fun assertPreviewUiState(
-        slots: Map<String, KeyguardQuickAffordanceSlotViewModel>,
+        slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
         expectedAffordanceNameBySlotId: Map<String, String?>,
     ) {
-        slots.forEach { (slotId, slotViewModel) ->
+        assertThat(slots).isNotNull()
+        slots?.forEach { (slotId, slotViewModel) ->
             val expectedAffordanceName = expectedAffordanceNameBySlotId[slotId]
             val actualAffordanceName =
                 slotViewModel.selectedQuickAffordances.firstOrNull()?.contentDescription
diff --git a/tests/src/com/android/customization/testing/TestCustomizationInjector.java b/tests/src/com/android/customization/testing/TestCustomizationInjector.java
index ed53751..f4a659c 100644
--- a/tests/src/com/android/customization/testing/TestCustomizationInjector.java
+++ b/tests/src/com/android/customization/testing/TestCustomizationInjector.java
@@ -12,15 +12,19 @@
 import com.android.customization.module.ThemesUserEventLogger;
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository;
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor;
-import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel;
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer;
 import com.android.systemui.shared.quickaffordance.data.content.KeyguardQuickAffordanceProviderClient;
 import com.android.systemui.shared.quickaffordance.data.content.KeyguardQuickAffordanceProviderClientImpl;
 import com.android.wallpaper.config.BaseFlags;
 import com.android.wallpaper.module.DrawableLayerResolver;
 import com.android.wallpaper.module.PackageStatusNotifier;
 import com.android.wallpaper.module.UserEventLogger;
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer;
 import com.android.wallpaper.testing.TestInjector;
 
+import java.util.HashMap;
+import java.util.Map;
+
 import kotlinx.coroutines.Dispatchers;
 
 /**
@@ -33,9 +37,9 @@
     private DrawableLayerResolver mDrawableLayerResolver;
     private UserEventLogger mUserEventLogger;
     private KeyguardQuickAffordancePickerInteractor mKeyguardQuickAffordancePickerInteractor;
-    private KeyguardQuickAffordancePickerViewModel.Factory
-            mKeyguardQuickAffordancePickerViewModelFactory;
     private BaseFlags mFlags;
+    private KeyguardQuickAffordanceProviderClient mKeyguardQuickAffordanceProviderClient;
+    private KeyguardQuickAffordanceSnapshotRestorer mKeyguardQuickAffordanceSnapshotRestorer;
 
     @Override
     public CustomizationPreferences getCustomizationPreferences(Context context) {
@@ -89,7 +93,8 @@
                     new KeyguardQuickAffordanceProviderClientImpl(context, Dispatchers.getIO());
             mKeyguardQuickAffordancePickerInteractor = new KeyguardQuickAffordancePickerInteractor(
                     new KeyguardQuickAffordancePickerRepository(client, Dispatchers.getIO()),
-                    client);
+                    client,
+                    () -> getKeyguardQuickAffordanceSnapshotRestorer(context));
         }
         return mKeyguardQuickAffordancePickerInteractor;
     }
@@ -102,4 +107,37 @@
 
         return mFlags;
     }
+
+    @Override
+    public Map<Integer, SnapshotRestorer> getSnapshotRestorers(Context context) {
+        final Map<Integer, SnapshotRestorer> restorers = new HashMap<>();
+        restorers.put(
+                KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER,
+                getKeyguardQuickAffordanceSnapshotRestorer(context));
+        return restorers;
+    }
+
+    /** Returns the {@link KeyguardQuickAffordanceProviderClient}. */
+    private KeyguardQuickAffordanceProviderClient getKeyguardQuickAffordancePickerProviderClient(
+            Context context) {
+        if (mKeyguardQuickAffordanceProviderClient == null) {
+            mKeyguardQuickAffordanceProviderClient =
+                    new KeyguardQuickAffordanceProviderClientImpl(context, Dispatchers.getIO());
+        }
+
+        return mKeyguardQuickAffordanceProviderClient;
+    }
+
+    private KeyguardQuickAffordanceSnapshotRestorer getKeyguardQuickAffordanceSnapshotRestorer(
+            Context context) {
+        if (mKeyguardQuickAffordanceSnapshotRestorer == null) {
+            mKeyguardQuickAffordanceSnapshotRestorer = new KeyguardQuickAffordanceSnapshotRestorer(
+                    getKeyguardQuickAffordancePickerInteractor(context),
+                    getKeyguardQuickAffordancePickerProviderClient(context));
+        }
+
+        return mKeyguardQuickAffordanceSnapshotRestorer;
+    }
+
+    private static final int KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER = 1;
 }