From 985c0beb8e47ad3f7334560adc40b10acc71a829 Mon Sep 17 00:00:00 2001
From: Austin Mroz <austin@comfy.org>
Date: Thu, 4 Sep 2025 13:31:49 -0500
Subject: [PATCH] Immediately display DOMWidgets when added

---
 src/components/selectionbar/SubgraphNodeWidget.vue | 14 ++++++++++++--
 src/lib/litegraph/src/subgraph/SubgraphNode.ts     |  3 ++-
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/components/selectionbar/SubgraphNodeWidget.vue b/src/components/selectionbar/SubgraphNodeWidget.vue
index 897b1f4fe6..c162ee556f 100644
--- a/src/components/selectionbar/SubgraphNodeWidget.vue
+++ b/src/components/selectionbar/SubgraphNodeWidget.vue
@@ -19,6 +19,7 @@
 <script setup lang="ts">
 import { ref, onMounted } from 'vue'
 import Button from 'primevue/button'
+import { useDomWidgetStore } from '@/stores/domWidgetStore'
 
 function hasWidget() {
   const node = props.node.data[2]
@@ -45,14 +46,23 @@ function onClick(e) {
   const node = props.node.data[2]
 
   console.log(isShown.value)
+  const { widgetStates } = useDomWidgetStore()
   if (!isShown.value) {
-    node.addProxyWidget(`${nodeId}`, widgetName)
+    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 {
     //FIXME: widget name collisions
     const index = node.widgets.findIndex((w) => w.name === widgetName)
     if (index < 0) throw new Error("Can't disable missing widget")
-    node.widgets.splice(index, 1)
+    const w = node.widgets.splice(index, 1)
+    if (widgetStates.has(w.id)) {
+      widgetStates.get(w.id).active = false
+    }
     isShown.value = false
   }
 }
diff --git a/src/lib/litegraph/src/subgraph/SubgraphNode.ts b/src/lib/litegraph/src/subgraph/SubgraphNode.ts
index 69a89bfa1d..35bc571c53 100644
--- a/src/lib/litegraph/src/subgraph/SubgraphNode.ts
+++ b/src/lib/litegraph/src/subgraph/SubgraphNode.ts
@@ -591,7 +591,7 @@ export class SubgraphNode extends LGraphNode implements BaseLGraph {
     this.properties.proxyWidgets ??= []
     //NOTE: This doesn't trigger onPropertyChanged
     this.properties.proxyWidgets.push([overlay.nodeId, overlay.widgetName])
-    this.addProxyFromOverlay({__proto__:overlay})
+    return this.addProxyFromOverlay({__proto__:overlay})
   }
   addProxyFromOverlay(overlay: Object) {
     overlay.graph = this.subgraph
@@ -635,5 +635,6 @@ export class SubgraphNode extends LGraphNode implements BaseLGraph {
     }))
     const w = new Proxy(overlay, handler)
     this.widgets.push(w)
+    return w
   }
 }