diff --git a/res/layout/customization_option_entry_app_shape.xml b/res/layout/customization_option_entry_app_shape.xml
deleted file mode 100644
index 965b0a1..0000000
--- a/res/layout/customization_option_entry_app_shape.xml
+++ /dev/null
@@ -1,47 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?><!--
-  ~ Copyright (C) 2024 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.
-  -->
-
-<androidx.constraintlayout.widget.ConstraintLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:app="http://schemas.android.com/apk/res-auto"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:paddingHorizontal="@dimen/customization_option_entry_horizontal_padding"
-    android:paddingVertical="@dimen/customization_option_entry_vertical_padding"
-    android:clickable="true">
-
-    <TextView
-        style="@style/CustomizationOptionEntryTitleTextStyle"
-        android:layout_width="0dp"
-        android:layout_height="wrap_content"
-        android:text="@string/preview_name_shape"
-        android:layout_marginEnd="@dimen/customization_option_entry_text_margin_end"
-        app:layout_constraintStart_toStartOf="parent"
-        app:layout_constraintEnd_toStartOf="@+id/option_entry_app_shape_icon"
-        app:layout_constraintBottom_toBottomOf="parent"
-        app:layout_constraintTop_toTopOf="parent"
-        app:layout_constraintVertical_chainStyle="packed" />
-
-    <FrameLayout
-        android:id="@+id/option_entry_app_shape_icon"
-        android:layout_width="@dimen/customization_option_entry_icon_size"
-        android:layout_height="@dimen/customization_option_entry_icon_size"
-        android:orientation="horizontal"
-        android:background="@drawable/customization_option_entry_icon_background"
-        app:layout_constraintEnd_toEndOf="parent"
-        app:layout_constraintTop_toTopOf="parent"
-        app:layout_constraintBottom_toBottomOf="parent" />
-</androidx.constraintlayout.widget.ConstraintLayout>
diff --git a/res/layout/customization_option_entry_app_grid.xml b/res/layout/customization_option_entry_app_shape_and_grid.xml
similarity index 100%
rename from res/layout/customization_option_entry_app_grid.xml
rename to res/layout/customization_option_entry_app_shape_and_grid.xml
diff --git a/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt
index f03bc9b..7217f61 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt
@@ -51,7 +51,17 @@
 
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                launch { viewModel.optionItems.collect { options -> adapter.setItems(options) } }
+                launch {
+                    viewModel.optionItems.collect { options ->
+                        adapter.setItems(options) {
+                            val indexToFocus =
+                                options.indexOfFirst { it.isSelected.value }.coerceAtLeast(0)
+                            (gridOptionList.layoutManager as LinearLayoutManager).scrollToPosition(
+                                indexToFocus
+                            )
+                        }
+                    }
+                }
             }
         }
     }
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
index d66e8a6..7ddafe0 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
@@ -96,7 +96,7 @@
 
         val optionShapeAndGrid =
             homeScreenCustomizationOptionEntries
-                .find { it.first == ThemePickerHomeCustomizationOption.APP_GRID }
+                .find { it.first == ThemePickerHomeCustomizationOption.APP_SHAPE_AND_GRID }
                 ?.second
         val optionShapeAndGridDescription =
             optionShapeAndGrid?.findViewById<TextView>(R.id.option_entry_app_grid_description)
@@ -222,7 +222,7 @@
             }
 
         customizationOptionFloatingSheetViewMap
-            ?.get(ThemePickerHomeCustomizationOption.APP_GRID)
+            ?.get(ThemePickerHomeCustomizationOption.APP_SHAPE_AND_GRID)
             ?.let {
                 ShapeAndGridFloatingSheetBinder.bind(
                     it,
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
index 6bddf4a..088a741 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -19,6 +19,7 @@
 import android.widget.Button
 import android.widget.FrameLayout
 import android.widget.Toolbar
+import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
@@ -54,12 +55,16 @@
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
-                    viewModel.keyguardQuickAffordancePickerViewModel2.onApply.collect { onApply ->
+                    viewModel.onApplyButtonClicked.collect { onApplyButtonClicked ->
                         applyButton.setOnClickListener {
-                            onApply?.invoke()?.let { viewModel.deselectOption() }
+                            onApplyButtonClicked?.invoke()?.let { viewModel.deselectOption() }
                         }
                     }
                 }
+
+                launch { viewModel.isOnApplyVisible.collect { applyButton.isVisible = it } }
+
+                launch { viewModel.isOnApplyEnabled.collect { applyButton.isEnabled = it } }
             }
         }
     }
diff --git a/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt b/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt
index 2ff62c9..176c757 100644
--- a/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt
+++ b/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt
@@ -44,8 +44,7 @@
 
     enum class ThemePickerHomeCustomizationOption : CustomizationOptionUtil.CustomizationOption {
         COLORS,
-        APP_GRID,
-        APP_SHAPE,
+        APP_SHAPE_AND_GRID,
         THEMED_ICONS,
     }
 
@@ -109,17 +108,9 @@
                             )
                     )
                     add(
-                        ThemePickerHomeCustomizationOption.APP_GRID to
+                        ThemePickerHomeCustomizationOption.APP_SHAPE_AND_GRID to
                             layoutInflater.inflate(
-                                R.layout.customization_option_entry_app_grid,
-                                optionContainer,
-                                false,
-                            )
-                    )
-                    add(
-                        ThemePickerHomeCustomizationOption.APP_SHAPE to
-                            layoutInflater.inflate(
-                                R.layout.customization_option_entry_app_shape,
+                                R.layout.customization_option_entry_app_shape_and_grid,
                                 optionContainer,
                                 false,
                             )
@@ -172,9 +163,9 @@
                     .also { bottomSheetContainer.addView(it) }
             )
             put(
-                ThemePickerHomeCustomizationOption.APP_GRID,
+                ThemePickerHomeCustomizationOption.APP_SHAPE_AND_GRID,
                 inflateFloatingSheet(
-                        ThemePickerHomeCustomizationOption.APP_GRID,
+                        ThemePickerHomeCustomizationOption.APP_SHAPE_AND_GRID,
                         bottomSheetContainer,
                         layoutInflater,
                     )
@@ -192,7 +183,8 @@
             ThemePickerLockCustomizationOption.CLOCK -> R.layout.floating_sheet_clock
             ThemePickerLockCustomizationOption.SHORTCUTS -> R.layout.floating_sheet_shortcut
             ThemePickerHomeCustomizationOption.COLORS -> R.layout.floating_sheet_colors
-            ThemePickerHomeCustomizationOption.APP_GRID -> R.layout.floating_sheet_shape_and_grid
+            ThemePickerHomeCustomizationOption.APP_SHAPE_AND_GRID ->
+                R.layout.floating_sheet_shape_and_grid
             else ->
                 throw IllegalStateException(
                     "Customization option $option does not have a bottom sheet view"
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
index 86cdd8a..658e435 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
@@ -83,12 +83,12 @@
                 started = SharingStarted.WhileSubscribed(),
                 initialValue = "",
             )
-    private val _selectedQuickAffordances = MutableStateFlow<Map<String, String>>(emptyMap())
-    val selectedQuickAffordances: Flow<Map<String, String>> =
-        _selectedQuickAffordances.asStateFlow()
+    private val _previewingQuickAffordances = MutableStateFlow<Map<String, String>>(emptyMap())
+    val previewingQuickAffordances: Flow<Map<String, String>> =
+        _previewingQuickAffordances.asStateFlow()
 
     fun resetPreview() {
-        _selectedQuickAffordances.tryEmit(emptyMap())
+        _previewingQuickAffordances.tryEmit(emptyMap())
         _selectedSlotId.tryEmit(SLOT_ID_BOTTOM_START)
     }
 
@@ -98,7 +98,7 @@
                 quickAffordanceInteractor.slots,
                 quickAffordanceInteractor.affordances,
                 quickAffordanceInteractor.selections,
-                selectedQuickAffordances,
+                previewingQuickAffordances,
                 selectedSlotId,
             ) { slots, affordances, selections, selectedQuickAffordances, selectedSlotId ->
                 slots.associate { slot ->
@@ -194,7 +194,7 @@
             val isNoneSelected =
                 combine(
                         selectedSlotId,
-                        selectedQuickAffordances,
+                        previewingQuickAffordances,
                         selectedAffordanceIds,
                     ) { selectedSlotId, selectedQuickAffordances, selectedAffordanceIds ->
                         selectedQuickAffordances[selectedSlotId]?.let {
@@ -214,10 +214,10 @@
                             if (!isSelected) {
                                 {
                                     val newMap =
-                                        _selectedQuickAffordances.value.toMutableMap().apply {
+                                        _previewingQuickAffordances.value.toMutableMap().apply {
                                             put(selectedSlotId, KEYGUARD_QUICK_AFFORDANCE_ID_NONE)
                                         }
-                                    _selectedQuickAffordances.tryEmit(newMap)
+                                    _previewingQuickAffordances.tryEmit(newMap)
                                 }
                             } else {
                                 null
@@ -230,7 +230,7 @@
                     val isSelectedFlow: StateFlow<Boolean> =
                         combine(
                                 selectedSlotId,
-                                selectedQuickAffordances,
+                                previewingQuickAffordances,
                                 selectedAffordanceIds,
                             ) { selectedSlotId, selectedQuickAffordances, selectedAffordanceIds ->
                                 selectedQuickAffordances[selectedSlotId]?.let {
@@ -255,10 +255,10 @@
                                     if (!isSelected) {
                                         {
                                             val newMap =
-                                                _selectedQuickAffordances.value
+                                                _previewingQuickAffordances.value
                                                     .toMutableMap()
                                                     .apply { put(selectedSlotId, affordance.id) }
-                                            _selectedQuickAffordances.tryEmit(newMap)
+                                            _previewingQuickAffordances.tryEmit(newMap)
                                         }
                                     } else {
                                         null
@@ -287,7 +287,7 @@
         }
 
     val onApply: Flow<(() -> Unit)?> =
-        selectedQuickAffordances.map {
+        previewingQuickAffordances.map {
             if (it.isEmpty()) {
                 null
             } else {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
index 0fdc0b0..2ce4747 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
@@ -49,8 +49,18 @@
     // The currently-set system grid option
     val selectedGridOption =
         interactor.selectedGridOption.filterNotNull().map { toOptionItemViewModel(it) }
+    private val _previewingGridOptionKey = MutableStateFlow<String?>(null)
+    // If the previewing key is null, use the currently-set system grid option
+    val previewingGridOptionKey =
+        combine(selectedGridOption, _previewingGridOptionKey) {
+            currentlySetGridOption,
+            previewingGridOptionKey ->
+            previewingGridOptionKey ?: currentlySetGridOption.key.value
+        }
 
-    private val overrideGridOptionKey = MutableStateFlow<String?>(null)
+    fun resetPreview() {
+        _previewingGridOptionKey.tryEmit(null)
+    }
 
     val optionItems: Flow<List<OptionItemViewModel<GridIconViewModel>>> =
         interactor.gridOptions.filterNotNull().map { gridOptions ->
@@ -58,21 +68,23 @@
         }
 
     val onApply: Flow<(() -> Unit)?> =
-        combine(selectedGridOption, overrideGridOptionKey) {
+        combine(selectedGridOption, _previewingGridOptionKey) {
             selectedGridOption,
-            overrideGridOptionKey ->
+            previewingGridOptionKey ->
             if (
-                overrideGridOptionKey == null ||
-                    overrideGridOptionKey == selectedGridOption.key.value
+                previewingGridOptionKey == null ||
+                    previewingGridOptionKey == selectedGridOption.key.value
             ) {
                 null
             } else {
-                { viewModelScope.launch { interactor.applySelectedOption(overrideGridOptionKey) } }
+                {
+                    viewModelScope.launch {
+                        interactor.applySelectedOption(previewingGridOptionKey)
+                    }
+                }
             }
         }
 
-    val isOnApplyEnabled: Flow<Boolean> = onApply.map { it != null }
-
     private fun toOptionItemViewModel(
         option: GridOptionModel
     ): OptionItemViewModel<GridIconViewModel> {
@@ -86,7 +98,7 @@
                     )
             )
         val isSelected =
-            overrideGridOptionKey
+            _previewingGridOptionKey
                 .map {
                     if (it == null) {
                         option.isCurrent
@@ -113,7 +125,7 @@
             onClicked =
                 isSelected.map {
                     if (!it) {
-                        { overrideGridOptionKey.value = option.key }
+                        { _previewingGridOptionKey.value = option.key }
                     } else {
                         null
                     }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index 4576e0c..0adff1b 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -25,8 +25,13 @@
 import dagger.assisted.AssistedInject
 import dagger.hilt.android.scopes.ViewModelScoped
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
 
 class ThemePickerCustomizationOptionsViewModel
 @AssistedInject
@@ -53,6 +58,8 @@
 
     override fun deselectOption(): Boolean {
         keyguardQuickAffordancePickerViewModel2.resetPreview()
+        shapeAndGridPickerViewModel.resetPreview()
+
         return defaultCustomizationOptionsViewModel.deselectOption()
     }
 
@@ -102,7 +109,7 @@
                 {
                     defaultCustomizationOptionsViewModel.selectOption(
                         ThemePickerCustomizationOptionUtil.ThemePickerHomeCustomizationOption
-                            .APP_GRID
+                            .APP_SHAPE_AND_GRID
                     )
                 }
             } else {
@@ -110,6 +117,24 @@
             }
         }
 
+    @OptIn(ExperimentalCoroutinesApi::class)
+    val onApplyButtonClicked =
+        selectedOption
+            .flatMapLatest {
+                when (it) {
+                    ThemePickerCustomizationOptionUtil.ThemePickerLockCustomizationOption
+                        .SHORTCUTS -> keyguardQuickAffordancePickerViewModel2.onApply
+                    ThemePickerCustomizationOptionUtil.ThemePickerHomeCustomizationOption
+                        .APP_SHAPE_AND_GRID -> shapeAndGridPickerViewModel.onApply
+                    else -> flow { emit(null) }
+                }
+            }
+            .stateIn(viewModelScope, SharingStarted.Eagerly, null)
+
+    val isOnApplyEnabled: Flow<Boolean> = onApplyButtonClicked.map { it != null }
+
+    val isOnApplyVisible: Flow<Boolean> = selectedOption.map { it != null }
+
     @ViewModelScoped
     @AssistedFactory
     interface Factory : CustomizationOptionsViewModelFactory {
diff --git a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
index 7d40b5b..eec7d5a 100644
--- a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
+++ b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
@@ -18,6 +18,7 @@
 
 import android.os.Bundle
 import android.os.Message
+import androidx.core.os.bundleOf
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
@@ -65,65 +66,86 @@
             )
         }
 
-        if (screen == Screen.LOCK_SCREEN) {
-            lifecycleOwner.lifecycleScope.launch {
-                lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                    launch {
-                        viewModel.selectedOption.collect {
-                            when (it) {
-                                ThemePickerLockCustomizationOption.SHORTCUTS ->
-                                    workspaceCallback.sendMessage(
-                                        MESSAGE_ID_START_CUSTOMIZING_QUICK_AFFORDANCES,
-                                        Bundle().apply {
-                                            putString(
-                                                KEY_INITIALLY_SELECTED_SLOT_ID,
-                                                SLOT_ID_BOTTOM_START,
-                                            )
-                                        }
-                                    )
-                                else ->
-                                    workspaceCallback.sendMessage(
-                                        MESSAGE_ID_DEFAULT_PREVIEW,
-                                        Bundle.EMPTY,
-                                    )
-                            }
-                        }
-                    }
-
-                    launch {
-                        viewModel.keyguardQuickAffordancePickerViewModel2.selectedSlotId.collect {
-                            workspaceCallback.sendMessage(
-                                MESSAGE_ID_SLOT_SELECTED,
-                                Bundle().apply { putString(KEY_SLOT_ID, it) },
-                            )
-                        }
-                    }
-
-                    launch {
-                        viewModel.keyguardQuickAffordancePickerViewModel2.selectedQuickAffordances
-                            .collect {
-                                it[SLOT_ID_BOTTOM_START]?.let {
-                                    workspaceCallback.sendMessage(
-                                        MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
-                                        Bundle().apply {
-                                            putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_START)
-                                            putString(KEY_QUICK_AFFORDANCE_ID, it)
-                                        },
-                                    )
-                                }
-                                it[SLOT_ID_BOTTOM_END]?.let {
-                                    workspaceCallback.sendMessage(
-                                        MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
-                                        Bundle().apply {
-                                            putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_END)
-                                            putString(KEY_QUICK_AFFORDANCE_ID, it)
-                                        },
-                                    )
+        when (screen) {
+            Screen.LOCK_SCREEN ->
+                lifecycleOwner.lifecycleScope.launch {
+                    lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                        launch {
+                            viewModel.selectedOption.collect {
+                                when (it) {
+                                    ThemePickerLockCustomizationOption.SHORTCUTS ->
+                                        workspaceCallback.sendMessage(
+                                            MESSAGE_ID_START_CUSTOMIZING_QUICK_AFFORDANCES,
+                                            Bundle().apply {
+                                                putString(
+                                                    KEY_INITIALLY_SELECTED_SLOT_ID,
+                                                    SLOT_ID_BOTTOM_START,
+                                                )
+                                            }
+                                        )
+                                    else ->
+                                        workspaceCallback.sendMessage(
+                                            MESSAGE_ID_DEFAULT_PREVIEW,
+                                            Bundle.EMPTY,
+                                        )
                                 }
                             }
+                        }
+
+                        launch {
+                            viewModel.keyguardQuickAffordancePickerViewModel2.selectedSlotId
+                                .collect {
+                                    workspaceCallback.sendMessage(
+                                        MESSAGE_ID_SLOT_SELECTED,
+                                        Bundle().apply { putString(KEY_SLOT_ID, it) },
+                                    )
+                                }
+                        }
+
+                        launch {
+                            viewModel.keyguardQuickAffordancePickerViewModel2
+                                .previewingQuickAffordances
+                                .collect {
+                                    it[SLOT_ID_BOTTOM_START]?.let {
+                                        workspaceCallback.sendMessage(
+                                            MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
+                                            Bundle().apply {
+                                                putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_START)
+                                                putString(KEY_QUICK_AFFORDANCE_ID, it)
+                                            },
+                                        )
+                                    }
+                                    it[SLOT_ID_BOTTOM_END]?.let {
+                                        workspaceCallback.sendMessage(
+                                            MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED,
+                                            Bundle().apply {
+                                                putString(KEY_SLOT_ID, SLOT_ID_BOTTOM_END)
+                                                putString(KEY_QUICK_AFFORDANCE_ID, it)
+                                            },
+                                        )
+                                    }
+                                }
+                        }
                     }
                 }
-            }
+            Screen.HOME_SCREEN ->
+                lifecycleOwner.lifecycleScope.launch {
+                    lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                        launch {
+                            viewModel.shapeAndGridPickerViewModel.previewingGridOptionKey.collect {
+                                workspaceCallback.sendMessage(
+                                    MESSAGE_ID_UPDATE_GRID,
+                                    bundleOf(KEY_GRID_NAME to it)
+                                )
+                            }
+                        }
+                    }
+                }
         }
     }
+
+    companion object {
+        const val MESSAGE_ID_UPDATE_GRID = 7414
+        const val KEY_GRID_NAME = "grid_name"
+    }
 }
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 19d74f6..b6f249e 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
@@ -117,19 +117,19 @@
     @Test
     fun selectedQuickAffordancesMapUpdates_whenClickingOnQuickAffordanceOptionsAndCallingResetPreview() =
         testScope.runTest {
-            val selectedQuickAffordances = collectLastValue(underTest.selectedQuickAffordances)
+            val previewingQuickAffordances = collectLastValue(underTest.previewingQuickAffordances)
 
             val tabs = collectLastValue(underTest.tabs)
             val quickAffordances = collectLastValue(underTest.quickAffordances)
 
             // Default selectedQuickAffordances is an empty map
-            assertThat(selectedQuickAffordances()).isEqualTo(emptyMap<String, String>())
+            assertThat(previewingQuickAffordances()).isEqualTo(emptyMap<String, String>())
 
             // Click on quick affordance 1 when selected slot ID is bottom_start
             val onClickAffordance1 =
                 collectLastValue(quickAffordances()?.get(1)?.onClicked ?: emptyFlow())
             onClickAffordance1()?.invoke()
-            assertThat(selectedQuickAffordances())
+            assertThat(previewingQuickAffordances())
                 .isEqualTo(
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
@@ -143,7 +143,7 @@
             val onClickAffordance2 =
                 collectLastValue(quickAffordances()?.get(2)?.onClicked ?: emptyFlow())
             onClickAffordance2()?.invoke()
-            assertThat(selectedQuickAffordances())
+            assertThat(previewingQuickAffordances())
                 .isEqualTo(
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
@@ -154,7 +154,7 @@
                 )
 
             underTest.resetPreview()
-            assertThat(selectedQuickAffordances()).isEqualTo(emptyMap<String, String>())
+            assertThat(previewingQuickAffordances()).isEqualTo(emptyMap<String, String>())
         }
 
     @Test