From 2d41aed0516d592b0d8cca70825a6d4108173b4f Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Wed, 26 Feb 2025 13:43:09 -0500
Subject: [PATCH] [BugFix] Properly update color button color on first
 selection (#2736)

---
 browser_tests/selectionToolbox.spec.ts    | 31 ++++++++++++++++++++---
 src/components/graph/SelectionToolbox.vue |  8 +++---
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/browser_tests/selectionToolbox.spec.ts b/browser_tests/selectionToolbox.spec.ts
index 18e64a626..8dd7bf9c7 100644
--- a/browser_tests/selectionToolbox.spec.ts
+++ b/browser_tests/selectionToolbox.spec.ts
@@ -4,6 +4,9 @@ import { comfyPageFixture } from './fixtures/ComfyPage'
 
 const test = comfyPageFixture
 
+const BLUE_COLOR = 'rgb(51, 51, 85)'
+const RED_COLOR = 'rgb(85, 51, 51)'
+
 test.describe('Selection Toolbox', () => {
   test.beforeEach(async ({ comfyPage }) => {
     await comfyPage.setSetting('Comfy.Canvas.SelectionToolbox', true)
@@ -150,10 +153,7 @@ test.describe('Selection Toolbox', () => {
       await redColorOption.click()
 
       // Button should now show the selected color
-      await expect(colorPickerButton).toHaveCSS(
-        'color',
-        'rgb(85, 51, 51)' // Red color, adjust if different
-      )
+      await expect(colorPickerButton).toHaveCSS('color', RED_COLOR)
     })
 
     test('color picker shows mixed state for differently colored selections', async ({
@@ -183,5 +183,28 @@ test.describe('Selection Toolbox', () => {
       )
       await expect(colorPickerButton).not.toHaveAttribute('color')
     })
+
+    test('color picker shows correct color when selecting pre-colored node', async ({
+      comfyPage
+    }) => {
+      // First color a node
+      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()
+
+      // Clear selection
+      await comfyPage.selectNodes(['KSampler'])
+
+      // Re-select the node
+      await comfyPage.selectNodes(['KSampler'])
+
+      // Color picker button should show the correct color
+      const colorPickerButton = comfyPage.page.locator(
+        '.selection-toolbox .pi-circle-fill'
+      )
+      await expect(colorPickerButton).toHaveCSS('color', BLUE_COLOR)
+    })
   })
 })
diff --git a/src/components/graph/SelectionToolbox.vue b/src/components/graph/SelectionToolbox.vue
index 864c254a3..641c52cd5 100644
--- a/src/components/graph/SelectionToolbox.vue
+++ b/src/components/graph/SelectionToolbox.vue
@@ -6,9 +6,9 @@
       content: 'p-0 flex flex-row'
     }"
   >
-    <ColorPickerButton v-if="nodeSelected || groupSelected" />
+    <ColorPickerButton v-show="nodeSelected || groupSelected" />
     <Button
-      v-if="nodeSelected"
+      v-show="nodeSelected"
       severity="secondary"
       text
       @click="
@@ -21,7 +21,7 @@
       </template>
     </Button>
     <Button
-      v-if="nodeSelected || groupSelected"
+      v-show="nodeSelected || groupSelected"
       severity="secondary"
       text
       icon="pi pi-thumbtack"
@@ -34,7 +34,7 @@
       @click="() => commandStore.execute('Comfy.Canvas.DeleteSelectedItems')"
     />
     <Button
-      v-if="isRefreshable"
+      v-show="isRefreshable"
       severity="info"
       text
       icon="pi pi-refresh"