From 3b3df250cdb89073cfa4b1e8734a0a4ffac9815c Mon Sep 17 00:00:00 2001
From: bymyself <abolkonsky.rem@gmail.com>
Date: Tue, 18 Feb 2025 09:39:43 -0700
Subject: [PATCH] Add refresh button to selecton toolbox (#2612)

---
 browser_tests/remoteWidgets.spec.ts        | 17 ++++++
 browser_tests/selectionToolbox.spec.ts     | 12 +++++
 src/components/graph/SelectionToolbox.vue  |  9 ++++
 src/composables/useRefreshableSelection.ts | 61 ++++++++++++++++++++++
 src/composables/widgets/useRemoteWidget.ts | 15 ++++--
 src/types/litegraph-augmentation.d.ts      |  5 ++
 6 files changed, 114 insertions(+), 5 deletions(-)
 create mode 100644 src/composables/useRefreshableSelection.ts

diff --git a/browser_tests/remoteWidgets.spec.ts b/browser_tests/remoteWidgets.spec.ts
index 285570fd9..972a8ffde 100644
--- a/browser_tests/remoteWidgets.spec.ts
+++ b/browser_tests/remoteWidgets.spec.ts
@@ -176,6 +176,23 @@ test.describe('Remote COMBO Widget', () => {
   })
 
   test.describe('Refresh Behavior', () => {
+    test('refresh button is visible in selection toolbar when node is selected', async ({
+      comfyPage
+    }) => {
+      await comfyPage.setSetting('Comfy.Canvas.SelectionToolbox', true)
+
+      const nodeName = 'Remote Widget Node'
+      await addRemoteWidgetNode(comfyPage, nodeName)
+      await waitForWidgetUpdate(comfyPage)
+
+      // Select remote widget node
+      await comfyPage.page.keyboard.press('Control+A')
+
+      await expect(
+        comfyPage.page.locator('.selection-toolbox .pi-refresh')
+      ).toBeVisible()
+    })
+
     test('refreshes options when TTL expires', async ({ comfyPage }) => {
       // Fulfill each request with a unique timestamp
       await comfyPage.page.route(
diff --git a/browser_tests/selectionToolbox.spec.ts b/browser_tests/selectionToolbox.spec.ts
index db47f9bd1..6bd1a87e3 100644
--- a/browser_tests/selectionToolbox.spec.ts
+++ b/browser_tests/selectionToolbox.spec.ts
@@ -57,4 +57,16 @@ test.describe('Selection Toolbox', () => {
       comfyPage.page.locator('.selection-overlay-container.show-border')
     ).not.toBeVisible()
   })
+
+  test('displays refresh button in toolbox when all nodes are selected', async ({
+    comfyPage
+  }) => {
+    // Select all nodes
+    await comfyPage.page.focus('canvas')
+    await comfyPage.page.keyboard.press('Control+A')
+
+    await expect(
+      comfyPage.page.locator('.selection-toolbox .pi-refresh')
+    ).toBeVisible()
+  })
 })
diff --git a/src/components/graph/SelectionToolbox.vue b/src/components/graph/SelectionToolbox.vue
index 9f2f3c39f..ead0fde8e 100644
--- a/src/components/graph/SelectionToolbox.vue
+++ b/src/components/graph/SelectionToolbox.vue
@@ -18,6 +18,13 @@
       icon="pi pi-trash"
       @click="() => commandStore.execute('Comfy.Canvas.DeleteSelectedItems')"
     />
+    <Button
+      v-if="isRefreshable"
+      severity="info"
+      text
+      icon="pi pi-refresh"
+      @click="refreshSelected"
+    />
   </Panel>
 </template>
 
@@ -25,9 +32,11 @@
 import Button from 'primevue/button'
 import Panel from 'primevue/panel'
 
+import { useRefreshableSelection } from '@/composables/useRefreshableSelection'
 import { useCommandStore } from '@/stores/commandStore'
 
 const commandStore = useCommandStore()
+const { isRefreshable, refreshSelected } = useRefreshableSelection()
 </script>
 
 <style scoped>
diff --git a/src/composables/useRefreshableSelection.ts b/src/composables/useRefreshableSelection.ts
new file mode 100644
index 000000000..8cce3029a
--- /dev/null
+++ b/src/composables/useRefreshableSelection.ts
@@ -0,0 +1,61 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+import type { IWidget } from '@comfyorg/litegraph'
+import { computed, ref, watchEffect } from 'vue'
+
+import { useCommandStore } from '@/stores/commandStore'
+import { useCanvasStore } from '@/stores/graphStore'
+
+interface RefreshableItem {
+  refresh: () => Promise<void> | void
+}
+
+type RefreshableWidget = IWidget & RefreshableItem
+
+const isLGraphNode = (item: unknown): item is LGraphNode => {
+  const name = item?.constructor?.name
+  return name === 'ComfyNode' || name === 'LGraphNode'
+}
+
+const isRefreshableWidget = (widget: IWidget): widget is RefreshableWidget =>
+  'refresh' in widget && typeof widget.refresh === 'function'
+
+/**
+ * Tracks selected nodes and their refreshable widgets
+ */
+export const useRefreshableSelection = () => {
+  const graphStore = useCanvasStore()
+  const commandStore = useCommandStore()
+  const selectedNodes = ref<LGraphNode[]>([])
+  const isAllNodesSelected = ref(false)
+
+  watchEffect(() => {
+    selectedNodes.value = graphStore.selectedItems.filter(isLGraphNode)
+    isAllNodesSelected.value =
+      graphStore.canvas?.graph?.nodes?.every((node) => !!node.selected) ?? false
+  })
+
+  const refreshableWidgets = computed(() =>
+    selectedNodes.value.flatMap(
+      (node) => node.widgets?.filter(isRefreshableWidget) ?? []
+    )
+  )
+
+  const isRefreshable = computed(
+    () => refreshableWidgets.value.length > 0 || isAllNodesSelected.value
+  )
+
+  async function refreshSelected() {
+    if (!isRefreshable.value) return
+
+    if (isAllNodesSelected.value) {
+      commandStore.execute('Comfy.RefreshNodeDefinitions')
+    } else {
+      await Promise.all(refreshableWidgets.value.map((item) => item.refresh()))
+    }
+  }
+
+  return {
+    isRefreshable,
+    refreshSelected
+  }
+}
diff --git a/src/composables/widgets/useRemoteWidget.ts b/src/composables/widgets/useRemoteWidget.ts
index 765605ddc..dd9b25459 100644
--- a/src/composables/widgets/useRemoteWidget.ts
+++ b/src/composables/widgets/useRemoteWidget.ts
@@ -78,6 +78,7 @@ export function useRemoteWidget<
   const isPermanent = refresh <= 0
   const cacheKey = createCacheKey(config)
   let isLoaded = false
+  let refreshQueued = false
 
   const setSuccess = (entry: CacheEntry<T>, data: T) => {
     entry.retryCount = 0
@@ -183,12 +184,15 @@ export function useRemoteWidget<
   /**
    * Getter of the remote property of the widget (e.g., options.values, value, etc.).
    * Starts the fetch process then returns the cached value immediately.
-   * @param onFulfilled - Optional callback to be called when the fetch is resolved.
    * @returns the most recent value of the widget.
    */
   function getValue(onFulfilled?: () => void) {
     fetchValue().then((data) => {
       if (isFirstLoad()) onFirstLoad(data)
+      if (refreshQueued && data !== defaultValue) {
+        onRefresh()
+        refreshQueued = false
+      }
       onFulfilled?.()
     })
     return getCachedValue() ?? defaultValue
@@ -197,22 +201,23 @@ export function useRemoteWidget<
   /**
    * Force the widget to refresh its value
    */
-  function refreshValue() {
+  widget.refresh = function () {
+    refreshQueued = true
     clearCachedValue()
-    getValue(onRefresh)
+    getValue()
   }
 
   /**
    * Add a refresh button to the node that, when clicked, will force the widget to refresh
    */
   function addRefreshButton() {
-    node.addWidget('button', 'refresh', 'refresh', refreshValue)
+    node.addWidget('button', 'refresh', 'refresh', widget.refresh)
   }
 
   return {
     getCachedValue,
     getValue,
-    refreshValue,
+    refreshValue: widget.refresh,
     addRefreshButton,
     getCacheEntry: () => dataCache.get(cacheKey),
 
diff --git a/src/types/litegraph-augmentation.d.ts b/src/types/litegraph-augmentation.d.ts
index bbf9c2a04..1c71d1ccd 100644
--- a/src/types/litegraph-augmentation.d.ts
+++ b/src/types/litegraph-augmentation.d.ts
@@ -22,6 +22,11 @@ declare module '@comfyorg/litegraph/dist/types/widgets' {
       index: number
     ) => Promise<unknown> | unknown
 
+    /**
+     * Refreshes the widget's value or options from its remote source.
+     */
+    refresh?: () => unknown
+
     /**
      * If the widget supports dynamic prompts, this will be set to true.
      * See extensions/core/dynamicPrompts.ts