diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index 7138add..2bf36be 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -83,6 +83,7 @@
         // Wallpaper quick switch section.
         sectionControllers.add(
                 new WallpaperQuickSwitchSectionController(
+                        screen,
                         wallpaperQuickSwitchViewModel,
                         lifecycleOwner,
                         sectionNavigationController));
diff --git a/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt b/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt
index 904d2f9..7bb3232 100644
--- a/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt
+++ b/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt
@@ -18,12 +18,13 @@
 package com.android.customization.picker.clock.domain.interactor
 
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
 import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
 
 /** Handles state restoration for clocks. */
 class ClocksSnapshotRestorer : SnapshotRestorer {
     override suspend fun setUpSnapshotRestorer(
-        updater: (RestorableSnapshot) -> Unit,
+        store: SnapshotStore,
     ): RestorableSnapshot {
         // TODO(b/262924055): implement as part of the clock settings screen.
         return RestorableSnapshot(mapOf())
diff --git a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
index cb2dbdc..2dfa38a 100644
--- a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
+++ b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
@@ -19,6 +19,7 @@
 
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
 import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
 
 /** Handles state restoration for the quick affordances system. */
@@ -27,16 +28,16 @@
     private val client: CustomizationProviderClient,
 ) : SnapshotRestorer {
 
-    private lateinit var snapshotUpdater: (RestorableSnapshot) -> Unit
+    private lateinit var snapshotStore: SnapshotStore
 
     suspend fun storeSnapshot() {
-        snapshotUpdater(snapshot())
+        snapshotStore.store(snapshot())
     }
 
     override suspend fun setUpSnapshotRestorer(
-        updater: (RestorableSnapshot) -> Unit,
+        store: SnapshotStore,
     ): RestorableSnapshot {
-        snapshotUpdater = updater
+        snapshotStore = store
         return snapshot()
     }
 
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 9a2a0af..4879fb0 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
@@ -24,6 +24,7 @@
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
@@ -69,7 +70,7 @@
                             interactor = underTest,
                             client = client,
                         )
-                        .apply { runBlocking { setUpSnapshotRestorer {} } }
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
                 },
             )
     }
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 cf09f9b..be118ee 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
@@ -35,6 +35,7 @@
 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.FakeSnapshotStore
 import com.android.wallpaper.testing.TestCurrentWallpaperInfoFactory
 import com.android.wallpaper.testing.TestInjector
 import com.android.wallpaper.testing.collectLastValue
@@ -90,7 +91,7 @@
                             interactor = quickAffordanceInteractor,
                             client = client,
                         )
-                        .apply { runBlocking { setUpSnapshotRestorer {} } }
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
                 },
             )
         val undoInteractor =