diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
index fd04580..32ee209 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
@@ -51,8 +51,8 @@
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.take
 
 class KeyguardQuickAffordancePickerViewModel2
 @AssistedInject
@@ -80,12 +80,31 @@
                 started = SharingStarted.WhileSubscribed(),
                 initialValue = "",
             )
-    private val _previewingQuickAffordances = MutableStateFlow<Map<String, String>>(emptyMap())
-    val previewingQuickAffordances: Flow<Map<String, String>> =
-        _previewingQuickAffordances.asStateFlow()
+    private val overridingQuickAffordances = MutableStateFlow<Map<String, String>>(emptyMap())
+    private val selectedQuickAffordancesGroupBySlotId =
+        quickAffordanceInteractor.selections.map {
+            it.groupBy { selectionModel -> selectionModel.slotId }
+        }
+
+    val previewingQuickAffordances =
+        combine(
+            quickAffordanceInteractor.slots,
+            overridingQuickAffordances,
+            selectedQuickAffordancesGroupBySlotId,
+        ) { slots, overridingQuickAffordances, selectedQuickAffordancesGroupBySlotId ->
+            slots.associate { slot ->
+                val selectedAffordanceId =
+                    overridingQuickAffordances[slot.id]
+                        ?: selectedQuickAffordancesGroupBySlotId[slot.id]
+                            ?.firstOrNull()
+                            ?.affordanceId
+                        ?: KEYGUARD_QUICK_AFFORDANCE_ID_NONE
+                slot.id to selectedAffordanceId
+            }
+        }
 
     fun resetPreview() {
-        _previewingQuickAffordances.tryEmit(emptyMap())
+        overridingQuickAffordances.tryEmit(emptyMap())
         _selectedSlotId.tryEmit(SLOT_ID_BOTTOM_START)
     }
 
@@ -94,21 +113,13 @@
         combine(
                 quickAffordanceInteractor.slots,
                 quickAffordanceInteractor.affordances,
-                quickAffordanceInteractor.selections,
                 previewingQuickAffordances,
                 selectedSlotId,
-            ) { slots, affordances, selections, selectedQuickAffordances, selectedSlotId ->
+            ) { slots, affordances, previewingQuickAffordances, selectedSlotId ->
                 slots.associate { slot ->
-                    val selectedAffordanceIds =
-                        selectedQuickAffordances[slot.id]?.let { setOf(it) }
-                            ?: selections
-                                .filter { selection -> selection.slotId == slot.id }
-                                .map { selection -> selection.affordanceId }
-                                .toSet()
+                    val selectedAffordanceId = previewingQuickAffordances[slot.id]
                     val selectedAffordances =
-                        affordances.filter { affordance ->
-                            selectedAffordanceIds.contains(affordance.id)
-                        }
+                        affordances.filter { affordance -> selectedAffordanceId == affordance.id }
 
                     val isSelected = selectedSlotId == slot.id
                     slot.id to
@@ -160,31 +171,15 @@
             }
         }
 
-    /**
-     * The set of IDs of the currently-selected affordances. These change with user selection of new
-     * or different affordances in the currently-selected slot or when slot selection changes.
-     */
-    private val selectedAffordanceIds: Flow<Set<String>> =
-        combine(quickAffordanceInteractor.selections, selectedSlotId) { selections, selectedSlotId
-                ->
-                selections
-                    .filter { selection -> selection.slotId == selectedSlotId }
-                    .map { selection -> selection.affordanceId }
-                    .toSet()
-            }
-            .shareIn(scope = viewModelScope, started = SharingStarted.WhileSubscribed(), replay = 1)
-
     /** The list of all available quick affordances for the selected slot. */
     val quickAffordances: Flow<List<OptionItemViewModel2<Icon>>> =
         quickAffordanceInteractor.affordances.map { affordances ->
             val isNoneSelected =
-                combine(selectedSlotId, previewingQuickAffordances, selectedAffordanceIds) {
+                combine(selectedSlotId, previewingQuickAffordances) {
                         selectedSlotId,
-                        selectedQuickAffordances,
-                        selectedAffordanceIds ->
-                        selectedQuickAffordances[selectedSlotId]?.let {
-                            it == KEYGUARD_QUICK_AFFORDANCE_ID_NONE
-                        } ?: selectedAffordanceIds.isEmpty()
+                        previewingQuickAffordances ->
+                        previewingQuickAffordances[selectedSlotId] ==
+                            KEYGUARD_QUICK_AFFORDANCE_ID_NONE
                     }
                     .stateIn(viewModelScope)
             listOf(
@@ -196,10 +191,10 @@
                             if (!isSelected) {
                                 {
                                     val newMap =
-                                        _previewingQuickAffordances.value.toMutableMap().apply {
+                                        overridingQuickAffordances.value.toMutableMap().apply {
                                             put(selectedSlotId, KEYGUARD_QUICK_AFFORDANCE_ID_NONE)
                                         }
-                                    _previewingQuickAffordances.tryEmit(newMap)
+                                    overridingQuickAffordances.tryEmit(newMap)
                                 }
                             } else {
                                 null
@@ -210,14 +205,10 @@
                 affordances.map { affordance ->
                     val affordanceIcon = getAffordanceIcon(affordance.iconResourceId)
                     val isSelectedFlow: StateFlow<Boolean> =
-                        combine(
+                        combine(selectedSlotId, previewingQuickAffordances) {
                                 selectedSlotId,
-                                previewingQuickAffordances,
-                                selectedAffordanceIds,
-                            ) { selectedSlotId, selectedQuickAffordances, selectedAffordanceIds ->
-                                selectedQuickAffordances[selectedSlotId]?.let {
-                                    it == affordance.id
-                                } ?: selectedAffordanceIds.contains(affordance.id)
+                                previewingQuickAffordances ->
+                                previewingQuickAffordances[selectedSlotId] == affordance.id
                             }
                             .stateIn(viewModelScope)
                     OptionItemViewModel2<Icon>(
@@ -235,10 +226,10 @@
                                     if (!isSelected) {
                                         {
                                             val newMap =
-                                                _previewingQuickAffordances.value
+                                                overridingQuickAffordances.value
                                                     .toMutableMap()
                                                     .apply { put(selectedSlotId, affordance.id) }
-                                            _previewingQuickAffordances.tryEmit(newMap)
+                                            overridingQuickAffordances.tryEmit(newMap)
                                         }
                                     } else {
                                         null
@@ -267,12 +258,20 @@
         }
 
     val onApply: Flow<(suspend () -> Unit)?> =
-        previewingQuickAffordances.map {
-            if (it.isEmpty()) {
-                null
-            } else {
+        combine(overridingQuickAffordances, selectedQuickAffordancesGroupBySlotId) {
+            overridingQuickAffordances,
+            selectedQuickAffordancesGroupBySlotId ->
+            // If all overridingQuickAffordances are same as the selected quick affordances, it is
+            // not yet edited
+            val isQuickAffordancesEdited =
+                (!overridingQuickAffordances.all { (slotId, overridingQuickAffordanceId) ->
+                    selectedQuickAffordancesGroupBySlotId[slotId]?.find {
+                        it.affordanceId == overridingQuickAffordanceId
+                    } != null
+                })
+            if (isQuickAffordancesEdited) {
                 {
-                    it.forEach { entry ->
+                    overridingQuickAffordances.forEach { entry ->
                         val slotId = entry.key
                         val affordanceId = entry.value
                         if (slotId == KEYGUARD_QUICK_AFFORDANCE_ID_NONE) {
@@ -283,9 +282,15 @@
                                 affordanceId = affordanceId,
                             )
                         }
+                        // Suspend until the next selectedQuickAffordancesGroupBySlotId update
+                        this.selectedQuickAffordancesGroupBySlotId.take(1).collect {
+                            return@collect
+                        }
                         logger.logShortcutApplied(shortcut = affordanceId, shortcutSlotId = slotId)
                     }
                 }
+            } else {
+                null
             }
         }
 
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt
index a7efc45..7053340 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt
@@ -26,6 +26,7 @@
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants
 import com.android.themepicker.R
 import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
@@ -123,7 +124,15 @@
             val quickAffordances = collectLastValue(underTest.quickAffordances)
 
             // Default selectedQuickAffordances is an empty map
-            assertThat(previewingQuickAffordances()).isEqualTo(emptyMap<String, String>())
+            assertThat(previewingQuickAffordances())
+                .isEqualTo(
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            KeyguardPreviewConstants.KEYGUARD_QUICK_AFFORDANCE_ID_NONE,
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            KeyguardPreviewConstants.KEYGUARD_QUICK_AFFORDANCE_ID_NONE,
+                    )
+                )
 
             // Click on quick affordance 1 when selected slot ID is bottom_start
             val onClickAffordance1 =
@@ -133,7 +142,9 @@
                 .isEqualTo(
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
-                            FakeCustomizationProviderClient.AFFORDANCE_1
+                            FakeCustomizationProviderClient.AFFORDANCE_1,
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            KeyguardPreviewConstants.KEYGUARD_QUICK_AFFORDANCE_ID_NONE,
                     )
                 )
 
@@ -154,7 +165,15 @@
                 )
 
             underTest.resetPreview()
-            assertThat(previewingQuickAffordances()).isEqualTo(emptyMap<String, String>())
+            assertThat(previewingQuickAffordances())
+                .isEqualTo(
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            KeyguardPreviewConstants.KEYGUARD_QUICK_AFFORDANCE_ID_NONE,
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            KeyguardPreviewConstants.KEYGUARD_QUICK_AFFORDANCE_ID_NONE,
+                    )
+                )
         }
 
     @Test