Fix shape on disconnect

src/lib/litegraph/src/canvas/ToInputFromIoNodeLink.ts

@@ -73,6 +73,7 @@ export class ToInputFromIoNodeLink implements RenderLink {
       if (inputNode && input)
         this.node._disconnectNodeInput(inputNode, input, existingLink)
       events.dispatch('input-moved', this)
+      fromSlot.events.dispatch('input-disconnected', { input: fromSlot })
     } else {
       // Creating a new link
       events.dispatch('link-created', newLink)
@@ -146,9 +147,17 @@ export class ToInputFromIoNodeLink implements RenderLink {
   }
   disconnect(): boolean {
     if (!this.existingLink) return false
-    const { input, inputNode } = this.existingLink.resolve(this.network)
+    const { input, inputNode, subgraphInput } = this.existingLink.resolve(
+      this.network
+    )
     if (!inputNode || !input) return false
     this.node._disconnectNodeInput(inputNode, input, this.existingLink)
+
+    if (subgraphInput)
+      subgraphInput.events.dispatch('input-disconnected', {
+        input: subgraphInput
+      })
+
     return true
   }
 }

src/lib/litegraph/src/subgraph/SubgraphNode.ts

@@ -764,7 +764,7 @@ export class SubgraphNode extends LGraphNode implements BaseLGraph {
       'input-connected',
       (e) => {
         this._invalidatePromotedViewsCache()
-        input.shape = e.detail.input.shape
+        input.shape = this.getSlotShape(subgraphInput, e.detail.input)
         if (!e.detail.widget || !e.detail.node) return
 
         // `SubgraphInput.connect()` dispatches before appending to `linkIds`,
@@ -807,6 +807,7 @@ export class SubgraphNode extends LGraphNode implements BaseLGraph {
       'input-disconnected',
       () => {
         this._invalidatePromotedViewsCache()
+        input.shape = this.getSlotShape(subgraphInput)
 
         // If links remain, rebind to the current representative.
         const connectedWidgets = subgraphInput.getConnectedWidgets()
@@ -895,19 +896,14 @@ export class SubgraphNode extends LGraphNode implements BaseLGraph {
 
     this.inputs.length = 0
     this.inputs.push(
-      ...this.subgraph.inputNode.slots.map((slot) => {
-        const shapes = slot.linkIds.map(
-          (id) => this.subgraph.links[id]?.resolve(this.subgraph)?.input?.shape
-        )
-        return Object.assign(
+      ...this.subgraph.inputNode.slots.map((slot) =>
+        Object.assign(
           new NodeInputSlot(
             {
               name: slot.name,
               localized_name: slot.localized_name,
               label: slot.label,
-              shape: shapes.every((shape) => shape === shapes[0])
-                ? shapes[0]
-                : undefined,
+              shape: this.getSlotShape(slot),
               type: slot.type,
               link: null
             },
@@ -917,7 +913,7 @@ export class SubgraphNode extends LGraphNode implements BaseLGraph {
             _subgraphSlot: slot
           }
         )
-      })
+      )
     )
 
     this.outputs.length = 0
@@ -1488,4 +1484,11 @@ export class SubgraphNode extends LGraphNode implements BaseLGraph {
 
     return clone
   }
+  getSlotShape(slot: SubgraphInput, extraInput?: INodeInputSlot) {
+    const shapes = slot.linkIds.map(
+      (id) => this.subgraph.links[id]?.resolve(this.subgraph)?.input?.shape
+    )
+    if (extraInput) shapes.push(extraInput.shape)
+    return shapes.every((shape) => shape === shapes[0]) ? shapes[0] : undefined
+  }
 }