Better subgraph cut/paste handling & tests

src/lib/ComfyGraph.ts

@@ -1,4 +1,4 @@
-import { LConnectionKind, LGraph, LGraphNode, type INodeSlot, type SlotIndex, LiteGraph, getStaticProperty, type LGraphAddNodeOptions, LGraphCanvas, type LGraphRemoveNodeOptions } from "@litegraph-ts/core";
+import { LConnectionKind, LGraph, LGraphNode, type INodeSlot, type SlotIndex, LiteGraph, getStaticProperty, type LGraphAddNodeOptions, LGraphCanvas, type LGraphRemoveNodeOptions, Subgraph, type LGraphAddNodeMode } from "@litegraph-ts/core";
 import GraphSync from "./GraphSync";
 import EventEmitter from "events";
 import type TypedEmitter from "typed-emitter";
@@ -41,8 +41,9 @@ export default class ComfyGraph extends LGraph {
     }
 
     override onNodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) {
-        if (options.subgraphs && options.subgraphs.length > 0)
-            return
+        // Don't add detached subgraphs
+        if (node.getRootGraph() == null || this._is_subgraph)
+            return;
 
         layoutState.nodeAdded(node, options)
 
@@ -112,7 +113,19 @@ export default class ComfyGraph extends LGraph {
                 const dragItemIDs = widgetNodesAdded.map(wn => get(layoutState).allItemsByNode[wn.id]?.dragItem?.id).filter(Boolean)
                 console.debug("[ComfyGraph] Group new widgets", dragItemIDs)
 
-                layoutState.groupItems(dragItemIDs, { title: node.title })
+                // Use the default node title instead of custom node title, in
+                // case node was cloned
+                const reg = LiteGraph.registered_node_types[node.type]
+
+                layoutState.groupItems(dragItemIDs, { title: reg.title })
+            }
+        }
+
+        // Handle subgraphs being attached
+        if (node.is(Subgraph)) {
+            console.error("ISSUBGRAPH")
+            for (const child of node.subgraph.iterateNodesInOrder()) {
+                this.onNodeAdded(child, options)
             }
         }
 
@@ -123,6 +136,13 @@ export default class ComfyGraph extends LGraph {
     override onNodeRemoved(node: LGraphNode, options: LGraphRemoveNodeOptions) {
         layoutState.nodeRemoved(node, options);
 
+        // Handle subgraphs being removed
+        if (node.is(Subgraph)) {
+            for (const child of node.subgraph.iterateNodesInOrder()) {
+                this.onNodeRemoved(child, options)
+            }
+        }
+
         // console.debug("Removed", node);
         this.eventBus.emit("nodeRemoved", node);
     }

src/lib/stores/layoutState.ts

@@ -291,7 +291,7 @@ const setNodeTitle = (arg: IDragItem, value: any) => {
     if (reg == null)
         return
 
-    if (value)
+    if (value && value !== reg.title)
         widget.node.title = `${reg.title} (${value})`
     else
         widget.node.title = reg.title
@@ -699,6 +699,20 @@ const store: Writable<LayoutState> = writable({
     }
 })
 
+function clear() {
+    store.set({
+        root: null,
+        allItems: {},
+        allItemsByNode: {},
+        isMenuOpen: false,
+        isConfiguring: true,
+        refreshPropsPanel: writable(0),
+        attrs: {
+            ...defaultWorkflowAttributes
+        }
+    })
+}
+
 function findDefaultContainerForInsertion(): ContainerLayout | null {
     const state = get(store);
 
@@ -741,6 +755,7 @@ function addContainer(parent: ContainerLayout | null, attrs: Partial<Attributes>
             ...attrs
         }
     }
+
     const entry: DragItemEntry = { dragItem, children: [], parent: null };
 
     if (state.allItemsByNode[dragItem.id] != null)
@@ -750,6 +765,7 @@ function addContainer(parent: ContainerLayout | null, attrs: Partial<Attributes>
     if (parent) {
         moveItem(dragItem, parent, index)
     }
+
     console.debug("[layoutState] addContainer", state)
     store.set(state)
     runOnChangedForWidgetDefaults(dragItem)
@@ -758,7 +774,7 @@ function addContainer(parent: ContainerLayout | null, attrs: Partial<Attributes>
 
 function addWidget(parent: ContainerLayout, node: ComfyWidgetNode, attrs: Partial<Attributes> = {}, index?: number): WidgetLayout {
     const state = get(store);
-    const widgetName = "Widget"
+    const widgetName = node.title || "Widget"
     const dragItem: WidgetLayout = {
         type: "widget",
         id: uuidv4(),
@@ -771,7 +787,7 @@ function addWidget(parent: ContainerLayout, node: ComfyWidgetNode, attrs: Partia
             ...attrs
         }
     }
-    const parentEntry = state.allItems[parent.id]
+
     const entry: DragItemEntry = { dragItem, children: [], parent: null };
 
     if (state.allItems[dragItem.id] != null)
@@ -821,6 +837,7 @@ function removeEntry(state: LayoutState, id: DragItemID) {
 
 function nodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) {
     const state = get(store)
+    console.error("NODEADDED", node.type, (node as any).svelteComponentType)
     if (state.isConfiguring)
         return;
 
@@ -834,7 +851,7 @@ function nodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) {
 
     const parent = findDefaultContainerForInsertion();
 
-    console.debug("[layoutState] nodeAdded", node)
+    console.debug("[layoutState] nodeAdded", node.id)
     if ("svelteComponentType" in node) {
         addWidget(parent, node as ComfyWidgetNode);
     }

src/lib/widgets/utils.ts

@@ -8,6 +8,9 @@ export function isNodeDisabled(node: LGraphNode): boolean {
         if (node.mode !== NodeMode.ALWAYS) {
             return true;
         }
+        if (node.graph == null) {
+            return true
+        }
         node = node.graph._subgraph_node;
     }
     return false;

src/tests/ComfyGraphTests.ts

@@ -0,0 +1,94 @@
+import { LGraph, LiteGraph, Subgraph, type SlotLayout } from "@litegraph-ts/core"
+import { Watch } from "@litegraph-ts/nodes-basic"
+import { expect } from 'vitest'
+import UnitTest from "./UnitTest"
+import ComfyGraph from "$lib/ComfyGraph";
+import ComfyPromptSerializer from "$lib/components/ComfyPromptSerializer";
+import { ComfyBackendNode } from "$lib/nodes/ComfyBackendNode";
+import ComfyGraphNode from "$lib/nodes/ComfyGraphNode";
+import { graphToGraphVis } from "$lib/utils";
+import layoutState from "$lib/stores/layoutState";
+import { ComfyNumberNode } from "$lib/nodes/widgets";
+import { get } from "svelte/store";
+
+export default class ComfyGraphTests extends UnitTest {
+    test__onNodeAdded__updatesLayoutState() {
+        const graph = new ComfyGraph();
+        layoutState.initDefaultLayout() // adds 3 containers
+
+        const state = get(layoutState)
+        expect(Object.keys(state.allItems)).toHaveLength(3)
+        expect(Object.keys(state.allItemsByNode)).toHaveLength(0)
+
+        const widget = LiteGraph.createNode(ComfyNumberNode);
+        graph.add(widget)
+
+        expect(Object.keys(state.allItems)).toHaveLength(4)
+        expect(Object.keys(state.allItemsByNode)).toHaveLength(1)
+        expect(state.allItemsByNode[widget.id]).toBeTruthy();
+
+        graph.add(widget)
+    }
+
+    test__correctSubgraphFactory() {
+        const graph = new ComfyGraph();
+        const subgraph = LiteGraph.createNode(Subgraph);
+        graph.add(subgraph)
+        expect(subgraph.graph).toBeInstanceOf(ComfyGraph)
+    }
+
+    test__onNodeAdded__handlesNodesAddedInSubgraphs() {
+        const graph = new ComfyGraph();
+        layoutState.initDefaultLayout()
+
+        const subgraph = LiteGraph.createNode(Subgraph);
+        graph.add(subgraph)
+
+        const state = get(layoutState)
+        expect(Object.keys(state.allItems)).toHaveLength(3)
+        expect(Object.keys(state.allItemsByNode)).toHaveLength(0)
+
+        const widget = LiteGraph.createNode(ComfyNumberNode);
+        subgraph.subgraph.add(widget)
+
+        expect(Object.keys(state.allItems)).toHaveLength(4)
+        expect(Object.keys(state.allItemsByNode)).toHaveLength(1)
+        expect(state.allItemsByNode[widget.id]).toBeTruthy();
+    }
+
+    test__onNodeAdded__handlesSubgraphsWithNodes() {
+        const graph = new ComfyGraph();
+        layoutState.initDefaultLayout()
+
+        const state = get(layoutState)
+        expect(Object.keys(state.allItems)).toHaveLength(3)
+        expect(Object.keys(state.allItemsByNode)).toHaveLength(0)
+
+        const subgraph = LiteGraph.createNode(Subgraph);
+        const widget = LiteGraph.createNode(ComfyNumberNode);
+        subgraph.subgraph.add(widget)
+        graph.add(subgraph)
+
+        expect(Object.keys(state.allItems)).toHaveLength(4)
+        expect(Object.keys(state.allItemsByNode)).toHaveLength(1)
+        expect(state.allItemsByNode[widget.id]).toBeTruthy();
+    }
+
+    test__onNodeRemoved__updatesLayoutState() {
+        const graph = new ComfyGraph();
+        layoutState.initDefaultLayout()
+
+        const widget = LiteGraph.createNode(ComfyNumberNode);
+        graph.add(widget)
+
+        const state = get(layoutState)
+        expect(Object.keys(state.allItems)).toHaveLength(4)
+        expect(Object.keys(state.allItemsByNode)).toHaveLength(1)
+        expect(state.allItemsByNode[widget.id]).toBeTruthy();
+
+        graph.remove(widget)
+
+        expect(Object.keys(state.allItems)).toHaveLength(3)
+        expect(Object.keys(state.allItemsByNode)).toHaveLength(0)
+    }
+}

src/tests/main.ts

@@ -1,3 +1,6 @@
+console.debug = (...msg) => {
+}
+
 import { vi, describe, it } from "vitest"
 import UnitTest from "./UnitTest"
 import * as testSuite from "./testSuite"

src/tests/testSuite.ts

@@ -1 +1,2 @@
 export { default as ComfyPromptSerializerTests } from "./ComfyPromptSerializerTests"
+export { default as ComfyGraphTests } from "./ComfyGraphTests"