From 197f33ffcd945dfed7d0a533f6ea22af354a84ff Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Fri, 2 May 2025 12:49:19 -0400
Subject: [PATCH] [Bug] Register dom widget when only node is added to graph
 (#3732)

---
 browser_tests/fixtures/ComfyPage.ts   |  6 ++++++
 browser_tests/tests/domWidget.spec.ts | 16 ++++++++++++++++
 src/scripts/domWidget.ts              | 11 +++++++++--
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/browser_tests/fixtures/ComfyPage.ts b/browser_tests/fixtures/ComfyPage.ts
index d0ac6b276..e27f440cd 100644
--- a/browser_tests/fixtures/ComfyPage.ts
+++ b/browser_tests/fixtures/ComfyPage.ts
@@ -924,6 +924,12 @@ export class ComfyPage {
       return window['app'].canvas.ds.convertOffsetToCanvas(pos)
     }, pos)
   }
+
+  /** Get number of DOM widgets on the canvas. */
+  async getDOMWidgetCount() {
+    return await this.page.locator('.dom-widget').count()
+  }
+
   async getNodeRefById(id: NodeId) {
     return new NodeReference(id, this)
   }
diff --git a/browser_tests/tests/domWidget.spec.ts b/browser_tests/tests/domWidget.spec.ts
index aa2fe89b8..4ff305e1a 100644
--- a/browser_tests/tests/domWidget.spec.ts
+++ b/browser_tests/tests/domWidget.spec.ts
@@ -31,4 +31,20 @@ test.describe('DOM Widget', () => {
     await comfyPage.nextFrame()
     await expect(comfyPage.canvas).toHaveScreenshot('focus-mode-on.png')
   })
+
+  // No DOM widget should be created by creation of interim LGraphNode objects.
+  test('Copy node with DOM widget by dragging + alt', async ({ comfyPage }) => {
+    const initialCount = await comfyPage.getDOMWidgetCount()
+
+    // TextEncodeNode1
+    await comfyPage.page.mouse.move(618, 191)
+    await comfyPage.page.keyboard.down('Alt')
+    await comfyPage.page.mouse.down()
+    await comfyPage.page.mouse.move(100, 100)
+    await comfyPage.page.mouse.up()
+    await comfyPage.page.keyboard.up('Alt')
+
+    const finalCount = await comfyPage.getDOMWidgetCount()
+    expect(finalCount).toBe(initialCount + 1)
+  })
 })
diff --git a/src/scripts/domWidget.ts b/src/scripts/domWidget.ts
index 933ef6f2b..1904bda5a 100644
--- a/src/scripts/domWidget.ts
+++ b/src/scripts/domWidget.ts
@@ -267,6 +267,15 @@ export const addWidget = <W extends BaseDOMWidget<object | string>>(
   widget: W
 ) => {
   node.addCustomWidget(widget)
+
+  if (node.graph) {
+    useDomWidgetStore().registerWidget(widget)
+  }
+
+  node.onAdded = useChainCallback(node.onAdded, () => {
+    useDomWidgetStore().registerWidget(widget)
+  })
+
   node.onRemoved = useChainCallback(node.onRemoved, () => {
     widget.onRemove?.()
   })
@@ -275,8 +284,6 @@ export const addWidget = <W extends BaseDOMWidget<object | string>>(
     widget.options.beforeResize?.call(widget, node)
     widget.options.afterResize?.call(widget, node)
   })
-
-  useDomWidgetStore().registerWidget(widget)
 }
 
 LGraphNode.prototype.addDOMWidget = function <