diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
index 4029dbe..9e2353a 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
@@ -26,6 +26,7 @@
 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
+import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import dagger.assisted.Assisted
@@ -33,8 +34,6 @@
 import dagger.assisted.AssistedInject
 import dagger.hilt.android.qualifiers.ApplicationContext
 import dagger.hilt.android.scopes.ViewModelScoped
-import kotlin.coroutines.resume
-import kotlinx.coroutines.CancellableContinuation
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -43,14 +42,15 @@
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.take
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.suspendCancellableCoroutine
 
 /** Models UI state for a color picker experience. */
 class ColorPickerViewModel2
 @AssistedInject
 constructor(
     @ApplicationContext context: Context,
+    private val colorUpdateViewModel: ColorUpdateViewModel,
     private val interactor: ColorPickerInteractor,
     private val logger: ThemesUserEventLogger,
     @Assisted private val viewModelScope: CoroutineScope,
@@ -60,7 +60,6 @@
     val previewingColorOption = overridingColorOption.asStateFlow()
 
     private val selectedColorTypeTabId = MutableStateFlow<ColorType?>(null)
-    private var onApplyContinuation: CancellableContinuation<Unit>? = null
 
     /** View-models for each color tab. */
     val colorTypeTabs: Flow<List<FloatingToolbarTabViewModel>> =
@@ -181,11 +180,9 @@
                 } else {
                     {
                         interactor.select(it)
-                        // Suspend until onApplyComplete is called, e.g. on configuration change
-                        suspendCancellableCoroutine { continuation: CancellableContinuation<Unit> ->
-                            onApplyContinuation?.cancel()
-                            onApplyContinuation = continuation
-                            continuation.invokeOnCancellation { onApplyContinuation = null }
+                        // Suspend until first color update
+                        colorUpdateViewModel.systemColorsUpdatedNoReplay.take(1).collect {
+                            return@collect
                         }
                         logger.logThemeColorApplied(
                             previewingColorOption.colorOption.sourceForLogging,
@@ -201,12 +198,6 @@
         overridingColorOption.value = null
     }
 
-    /** Resumes the onApply function if apply is in progress, otherwise no-op */
-    fun onApplyComplete() {
-        onApplyContinuation?.resume(Unit)
-        onApplyContinuation = null
-    }
-
     /** The list of all available color options for the selected Color Type. */
     val colorOptions: Flow<List<OptionItemViewModel2<ColorOptionIconViewModel>>> =
         combine(allColorOptions, selectedColorTypeTabId) {
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 4b79739..cb6057d 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
@@ -159,12 +159,6 @@
                         }
 
                         launch {
-                            colorUpdateViewModel.systemColorsUpdated.collect {
-                                viewModel.colorPickerViewModel2.onApplyComplete()
-                            }
-                        }
-
-                        launch {
                             combine(
                                     viewModel.colorPickerViewModel2.previewingColorOption,
                                     viewModel.darkModeViewModel.overridingIsDarkMode,
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
index e4a4884..9827a56 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
@@ -28,12 +28,14 @@
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.systemui.monet.Style
+import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
+import dagger.hilt.android.internal.lifecycle.RetainedLifecycleImpl
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.launch
@@ -58,6 +60,7 @@
     private lateinit var repository: FakeColorPickerRepository
     private lateinit var interactor: ColorPickerInteractor
     private lateinit var store: FakeSnapshotStore
+    private lateinit var colorUpdateViewModel: ColorUpdateViewModel
 
     private lateinit var context: Context
     private lateinit var testScope: TestScope
@@ -80,11 +83,14 @@
                     },
             )
 
+        colorUpdateViewModel = ColorUpdateViewModel(context, RetainedLifecycleImpl())
+
         underTest =
             ColorPickerViewModel2(
                 context = context,
                 interactor = interactor,
                 logger = logger,
+                colorUpdateViewModel = colorUpdateViewModel,
                 viewModelScope = testScope.backgroundScope,
             )
 
@@ -112,7 +118,7 @@
 
             assertThat(job.isActive).isTrue()
 
-            underTest.onApplyComplete()
+            colorUpdateViewModel.updateColors()
 
             assertThat(job.isActive).isFalse()
         }
@@ -243,7 +249,7 @@
     private suspend fun TestScope.applySelectedColorOption() {
         val onApply = collectLastValue(underTest.onApply)()
         testScope.launch { onApply?.invoke() }
-        underTest.onApplyComplete()
+        colorUpdateViewModel.updateColors()
     }
 
     /**