From 21e0caa1b1d6c22b8c34d17d3dacf86a648c2cd1 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Fri, 28 Mar 2025 21:00:43 -0400
Subject: [PATCH] [Bug] Fix undo of colorization via selection toolbox (#3267)

---
 browser_tests/tests/selectionToolbox.spec.ts  | 19 +++++++++++++++++++
 .../selectionToolbox/ColorPickerButton.vue    |  3 +++
 2 files changed, 22 insertions(+)

diff --git a/browser_tests/tests/selectionToolbox.spec.ts b/browser_tests/tests/selectionToolbox.spec.ts
index 6aa891fa08..35f3a7eb58 100644
--- a/browser_tests/tests/selectionToolbox.spec.ts
+++ b/browser_tests/tests/selectionToolbox.spec.ts
@@ -246,5 +246,24 @@ test.describe('Selection Toolbox', () => {
       )
       await expect(colorPickerButton).toHaveCSS('color', BLUE_COLOR)
     })
+
+    test('colorization via color picker can be undone', async ({
+      comfyPage
+    }) => {
+      // Select a node and color it
+      await comfyPage.selectNodes(['KSampler'])
+      await comfyPage.page.locator('.selection-toolbox .pi-circle-fill').click()
+      await comfyPage.page
+        .locator('.color-picker-container i[data-testid="blue"]')
+        .click()
+
+      // Undo the colorization
+      await comfyPage.page.keyboard.press('Control+Z')
+      await comfyPage.nextFrame()
+
+      // Node should be uncolored again
+      const selectedNode = (await comfyPage.getNodeRefsByTitle('KSampler'))[0]
+      expect(await selectedNode.getProperty('color')).toBeUndefined()
+    })
   })
 })
diff --git a/src/components/graph/selectionToolbox/ColorPickerButton.vue b/src/components/graph/selectionToolbox/ColorPickerButton.vue
index 2ce1556cb2..cf93f6cc22 100644
--- a/src/components/graph/selectionToolbox/ColorPickerButton.vue
+++ b/src/components/graph/selectionToolbox/ColorPickerButton.vue
@@ -47,6 +47,7 @@ import { computed, ref, watch } from 'vue'
 import { useI18n } from 'vue-i18n'
 
 import { useCanvasStore } from '@/stores/graphStore'
+import { useWorkflowStore } from '@/stores/workflowStore'
 import { useColorPaletteStore } from '@/stores/workspace/colorPaletteStore'
 import { adjustColor } from '@/utils/colorUtil'
 import { getItemsColorOption } from '@/utils/litegraphUtil'
@@ -54,6 +55,7 @@ import { getItemsColorOption } from '@/utils/litegraphUtil'
 const { t } = useI18n()
 const canvasStore = useCanvasStore()
 const colorPaletteStore = useColorPaletteStore()
+const workflowStore = useWorkflowStore()
 const isLightTheme = computed(
   () => colorPaletteStore.completedActivePalette.light_theme
 )
@@ -108,6 +110,7 @@ const applyColor = (colorOption: ColorOption | null) => {
   canvasStore.canvas?.setDirty(true, true)
   currentColorOption.value = canvasColorOption
   showColorPicker.value = false
+  workflowStore.activeWorkflow?.changeTracker.checkState()
 }
 
 const currentColorOption = ref<CanvasColorOption | null>(null)