diff --git a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayApplier.java b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayApplier.java
index d61848c95a09d..e0ff88bdbfcea 100644
--- a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayApplier.java
+++ b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayApplier.java
@@ -75,6 +75,10 @@ public class ThemeOverlayApplier implements Dumpable {
 
     static final String COLOR_SOURCE_PRESET = "preset";
 
+    static final String COLOR_SOURCE_HOME = "home_wallpaper";
+
+    static final String COLOR_SOURCE_LOCK = "lock_wallpaper";
+
     static final String TIMESTAMP_FIELD = "_applied_timestamp";
 
     @VisibleForTesting
diff --git a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
index c23e0b479c227..428921ee34ecd 100644
--- a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
+++ b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
@@ -16,6 +16,8 @@
 package com.android.systemui.theme;
 
 import static com.android.systemui.keyguard.WakefulnessLifecycle.WAKEFULNESS_ASLEEP;
+import static com.android.systemui.theme.ThemeOverlayApplier.COLOR_SOURCE_HOME;
+import static com.android.systemui.theme.ThemeOverlayApplier.COLOR_SOURCE_LOCK;
 import static com.android.systemui.theme.ThemeOverlayApplier.COLOR_SOURCE_PRESET;
 import static com.android.systemui.theme.ThemeOverlayApplier.OVERLAY_CATEGORY_ACCENT_COLOR;
 import static com.android.systemui.theme.ThemeOverlayApplier.OVERLAY_CATEGORY_SYSTEM_PALETTE;
@@ -200,36 +202,37 @@ public class ThemeOverlayController extends SystemUI implements Dumpable {
         String overlayPackageJson = mSecureSettings.getStringForUser(
                 Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES,
                 currentUser);
-        boolean isDestinationBoth = mWallpaperManager.getWallpaperId(
-                WallpaperManager.FLAG_LOCK) < 0;
-        if (!TextUtils.isEmpty(overlayPackageJson)) {
-            try {
-                JSONObject jsonObject = new JSONObject(overlayPackageJson);
-                if (!COLOR_SOURCE_PRESET.equals(jsonObject.optString(OVERLAY_COLOR_SOURCE))
-                        && ((flags & latestWallpaperType) != 0)) {
-                    mSkipSettingChange = true;
-                    if (jsonObject.has(OVERLAY_CATEGORY_ACCENT_COLOR) || jsonObject.has(
-                            OVERLAY_CATEGORY_SYSTEM_PALETTE)) {
-                        jsonObject.remove(OVERLAY_CATEGORY_ACCENT_COLOR);
-                        jsonObject.remove(OVERLAY_CATEGORY_SYSTEM_PALETTE);
-                        jsonObject.remove(OVERLAY_COLOR_SOURCE);
-                        jsonObject.remove(OVERLAY_COLOR_INDEX);
-                    }
-                    // Keep color_both value because users can change either or both home and
-                    // lock screen wallpapers.
-                    jsonObject.put(OVERLAY_COLOR_BOTH, isDestinationBoth ? "1" : "0");
-
-                    jsonObject.put(TIMESTAMP_FIELD, System.currentTimeMillis());
-                    if (DEBUG) {
-                        Log.d(TAG, "Updating theme setting from "
-                                + overlayPackageJson + " to " + jsonObject.toString());
-                    }
-                    mSecureSettings.putString(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES,
-                            jsonObject.toString());
+        boolean isDestinationBoth = (flags == (WallpaperManager.FLAG_SYSTEM
+                | WallpaperManager.FLAG_LOCK));
+        try {
+            JSONObject jsonObject = (overlayPackageJson == null) ? new JSONObject()
+                    : new JSONObject(overlayPackageJson);
+            if (!COLOR_SOURCE_PRESET.equals(jsonObject.optString(OVERLAY_COLOR_SOURCE))
+                    && ((flags & latestWallpaperType) != 0)) {
+                mSkipSettingChange = true;
+                if (jsonObject.has(OVERLAY_CATEGORY_ACCENT_COLOR) || jsonObject.has(
+                        OVERLAY_CATEGORY_SYSTEM_PALETTE)) {
+                    jsonObject.remove(OVERLAY_CATEGORY_ACCENT_COLOR);
+                    jsonObject.remove(OVERLAY_CATEGORY_SYSTEM_PALETTE);
+                    jsonObject.remove(OVERLAY_COLOR_INDEX);
                 }
-            } catch (JSONException e) {
-                Log.i(TAG, "Failed to parse THEME_CUSTOMIZATION_OVERLAY_PACKAGES.", e);
+                // Keep color_both value because users can change either or both home and
+                // lock screen wallpapers.
+                jsonObject.put(OVERLAY_COLOR_BOTH, isDestinationBoth ? "1" : "0");
+
+                jsonObject.put(OVERLAY_COLOR_SOURCE,
+                        (flags == WallpaperManager.FLAG_LOCK) ? COLOR_SOURCE_LOCK
+                                : COLOR_SOURCE_HOME);
+                jsonObject.put(TIMESTAMP_FIELD, System.currentTimeMillis());
+                if (DEBUG) {
+                    Log.d(TAG, "Updating theme setting from "
+                            + overlayPackageJson + " to " + jsonObject.toString());
+                }
+                mSecureSettings.putString(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES,
+                        jsonObject.toString());
             }
+        } catch (JSONException e) {
+            Log.i(TAG, "Failed to parse THEME_CUSTOMIZATION_OVERLAY_PACKAGES.", e);
         }
         reevaluateSystemTheme(false /* forceReload */);
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
index 3d2a0f1720750..3cf92126c5c60 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
@@ -232,9 +232,9 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
         verify(mSecureSettings).putString(
                 eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), updatedSetting.capture());
 
-        assertThat(updatedSetting.getValue().contains("android.theme.customization.system_palette"))
+        assertThat(updatedSetting.getValue().contains("android.theme.customization.accent_color"))
                 .isFalse();
-        assertThat(updatedSetting.getValue().contains("android.theme.customization.color_source"))
+        assertThat(updatedSetting.getValue().contains("android.theme.customization.system_palette"))
                 .isFalse();
         assertThat(updatedSetting.getValue().contains("android.theme.customization.color_index"))
                 .isFalse();
@@ -289,7 +289,8 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
                 .thenReturn(jsonString);
         when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(-1);
 
-        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+        mColorsListener.getValue().onColorsChanged(mainColors,
+                WallpaperManager.FLAG_SYSTEM | WallpaperManager.FLAG_LOCK);
 
         ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
         verify(mSecureSettings).putString(
@@ -302,6 +303,60 @@ public class ThemeOverlayControllerTest extends SysuiTestCase {
                 .applyCurrentUserOverlays(any(), any(), anyInt(), any());
     }
 
+    @Test
+    public void onWallpaperColorsChanged_changeLockWallpaper() {
+        // Should ask for a new theme when wallpaper colors change
+        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
+                Color.valueOf(Color.BLUE), null);
+        String jsonString =
+                "{\"android.theme.customization.system_palette\":\"override.package.name\","
+                        + "\"android.theme.customization.color_source\":\"home_wallpaper\","
+                        + "\"android.theme.customization.color_index\":\"2\"}";
+        when(mSecureSettings.getStringForUser(
+                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
+                .thenReturn(jsonString);
+        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(1);
+
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_LOCK);
+
+        ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
+        verify(mSecureSettings).putString(
+                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), updatedSetting.capture());
+        assertThat(updatedSetting.getValue().contains(
+                "android.theme.customization.color_source\":\"lock_wallpaper")).isTrue();
+        assertThat(updatedSetting.getValue().contains("android.theme.customization.color_index"))
+                .isFalse();
+        verify(mThemeOverlayApplier)
+                .applyCurrentUserOverlays(any(), any(), anyInt(), any());
+    }
+
+    @Test
+    public void onWallpaperColorsChanged_changeHomeWallpaper() {
+        // Should ask for a new theme when wallpaper colors change
+        WallpaperColors mainColors = new WallpaperColors(Color.valueOf(Color.RED),
+                Color.valueOf(Color.BLUE), null);
+        String jsonString =
+                "{\"android.theme.customization.system_palette\":\"override.package.name\","
+                        + "\"android.theme.customization.color_source\":\"lock_wallpaper\","
+                        + "\"android.theme.customization.color_index\":\"2\"}";
+        when(mSecureSettings.getStringForUser(
+                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), anyInt()))
+                .thenReturn(jsonString);
+        when(mWallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK)).thenReturn(-1);
+
+        mColorsListener.getValue().onColorsChanged(mainColors, WallpaperManager.FLAG_SYSTEM);
+
+        ArgumentCaptor<String> updatedSetting = ArgumentCaptor.forClass(String.class);
+        verify(mSecureSettings).putString(
+                eq(Settings.Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES), updatedSetting.capture());
+        assertThat(updatedSetting.getValue().contains(
+                "android.theme.customization.color_source\":\"home_wallpaper")).isTrue();
+        assertThat(updatedSetting.getValue().contains("android.theme.customization.color_index"))
+                .isFalse();
+        verify(mThemeOverlayApplier)
+                .applyCurrentUserOverlays(any(), any(), anyInt(), any());
+    }
+
     @Test
     public void onWallpaperColorsChanged_ResetThemeWhenFromLatestWallpaper() {
         // Should ask for a new theme when the colors of the last applied wallpaper change