From bb25301e367587dbcb0ef05188c946646e2c320a Mon Sep 17 00:00:00 2001
From: Austin Mroz <austin@comfy.org>
Date: Wed, 10 Sep 2025 15:37:31 -0500
Subject: [PATCH] Update widget state after visibility toggle

---
 src/components/selectionbar/SubgraphNode.vue  | 34 +++++++++++++-
 .../selectionbar/SubgraphNodeWidget.vue       | 47 ++-----------------
 2 files changed, 36 insertions(+), 45 deletions(-)

diff --git a/src/components/selectionbar/SubgraphNode.vue b/src/components/selectionbar/SubgraphNode.vue
index 4945470e5..3da178acd 100644
--- a/src/components/selectionbar/SubgraphNode.vue
+++ b/src/components/selectionbar/SubgraphNode.vue
@@ -7,6 +7,7 @@ import draggable from 'vuedraggable'
 
 import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'
 import SubgraphNodeWidget from '@/components/selectionbar/SubgraphNodeWidget.vue'
+import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { useCanvasStore } from '@/stores/graphStore'
 const { t } = useI18n()
 
@@ -27,6 +28,7 @@ const activeWidgets = computed({
   get() {
     triggerUpdate.value
     const node = activeNode.value
+    if (!node) return []
     const pw = node.properties.proxyWidgets ?? []
     return pw.map(([id, name]) => {
       const wNode = node.subgraph._nodes_by_id[id]
@@ -43,6 +45,32 @@ const activeWidgets = computed({
     canvasStore.canvas.setDirty(true)
   }
 })
+function toggleVisibility(nodeId, widgetName, isShown) {
+  const node = activeNode.value
+  const { widgetStates } = useDomWidgetStore()
+  if (!isShown) {
+    const w = node.addProxyWidget(`${nodeId}`, widgetName)
+    if (widgetStates.has(w.id)) {
+      const widgetState = widgetStates.get(w.id)
+      widgetState.active = true
+      widgetState.widget = w
+    }
+  } else {
+    const index = node.widgets.findIndex((w) => w.name === widgetName)
+    if (index < 0) throw new Error("Can't disable missing widget")
+    const [w] = node.widgets.splice(index, 1)
+    if (widgetStates.has(w.id)) {
+      widgetStates.get(w.id).active = false
+    }
+    const { properties } = node
+    properties.proxyWidgets = properties.proxyWidgets.filter((p) => {
+    return p[1] !== widgetName
+      //NOTE: intentional loose as nodeId is often string/int
+      || p[0] != nodeId})
+  }
+  triggerUpdate.value++
+  useCanvasStore().canvas.setDirty(true)
+}
 
 const candidateWidgets = computed(() =>{
   const node = canvasStore.selectedItems[0] ?? {}
@@ -78,13 +106,15 @@ const candidateWidgets = computed(() =>{
             @end="drag=false"
             item-key="id">
         <template #item="{element}">
-          <SubgraphNodeWidget :item="element" :node="activeNode" :draggable="true"/>
+          <SubgraphNodeWidget :item="element" :node="activeNode" :isShown="true"
+            :toggleVisibility="toggleVisibility"/>
         </template>
         </draggable>
         </div>
         <div class="widgets-section">
           <div v-for="element in candidateWidgets" class="widget-container">
-            <SubgraphNodeWidget :item="element" :node="activeNode"/>
+            <SubgraphNodeWidget :item="element" :node="activeNode"
+            :toggleVisibility="toggleVisibility"/>
           </div>
         </div>
       </template>
diff --git a/src/components/selectionbar/SubgraphNodeWidget.vue b/src/components/selectionbar/SubgraphNodeWidget.vue
index 7577ce5f6..3efe9917c 100644
--- a/src/components/selectionbar/SubgraphNodeWidget.vue
+++ b/src/components/selectionbar/SubgraphNodeWidget.vue
@@ -1,59 +1,20 @@
 <script setup lang="ts">
 import { ref, onMounted } from 'vue'
 import Button from 'primevue/button'
-import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { useCanvasStore } from '@/stores/graphStore'
 
-function hasWidget() {
-  return props.node.widgets.some((w) => w.name === props.item[1].name)
-}
-
-let isShown = ref(false)
-
 import TreeExplorerTreeNode from '@/components/common/TreeExplorerTreeNode.vue'
 import { RenderedTreeExplorerNode } from '@/types/treeExplorerTypes'
 
 const props = defineProps<{
   item: [unknown, unknown],
-  node: unknown
-  draggable?: boolean
+  node: unknown,
+  isShown?: boolean,
+  toggleVisibility
 }>()
 
-onMounted(() => {
-  isShown.value = hasWidget()
-})
-
 function onClick(e) {
-  //props.node?.onToggle()
-  const nodeId = props.item[0].id
-  const widgetName = props.item[1].name
-  const node = props.node
-
-  const { widgetStates } = useDomWidgetStore()
-  if (!isShown.value) {
-    const w = node.addProxyWidget(`${nodeId}`, widgetName)
-    if (widgetStates.has(w.id)) {
-      const widgetState = widgetStates.get(w.id)
-      widgetState.active = true
-      widgetState.widget = w
-    }
-    isShown.value = true
-  } else {
-    const index = node.widgets.findIndex((w) => w.name === widgetName)
-    if (index < 0) throw new Error("Can't disable missing widget")
-    const [w] = node.widgets.splice(index, 1)
-    if (widgetStates.has(w.id)) {
-      widgetStates.get(w.id).active = false
-    }
-    const { properties } = node
-    properties.proxyWidgets = properties.proxyWidgets.filter((p) => {
-    return p[1] !== widgetName
-      //NOTE: intentional loose as nodeId is often string/int
-      || p[0] != nodeId})
-
-    isShown.value = false
-  }
-  useCanvasStore().canvas.setDirty(true)
+  props.toggleVisibility(`${props.item[0].id}`, props.item[1].name, props.isShown)
 }
 </script>
 <template>