feat: centralize subgraph io mutation paths

Amp-Thread-ID: https://ampcode.com/threads/T-019c98a5-85ea-75ec-9594-aea357107cf6
Co-authored-by: Amp <amp@ampcode.com>

src/lib/litegraph/src/LGraph.ts

@@ -1622,6 +1622,27 @@ export class LGraph
     link.disconnect(this, keepReroutes)
   }
 
+  private finalizeConnectedLink(link: LLink): void {
+    const reroutes = LLink.getReroutes(this, link)
+    for (const reroute of reroutes) {
+      reroute.linkIds.add(link.id)
+      if (reroute.floating) reroute.floating = undefined
+      reroute._dragging = undefined
+    }
+
+    const lastReroute = reroutes.at(-1)
+    if (lastReroute) {
+      for (const floatingLinkId of lastReroute.floatingLinkIds) {
+        const floatingLink = this.floatingLinks.get(floatingLinkId)
+        if (floatingLink?.parentId === lastReroute.id) {
+          this.removeFloatingLink(floatingLink)
+        }
+      }
+    }
+
+    this._version++
+  }
+
   connectSlots(
     sourceNode: LGraphNode,
     outputIndex: number,
@@ -1674,27 +1695,112 @@ export class LGraph
       })
     }
 
-    // Reroutes
-    const reroutes = LLink.getReroutes(this, link)
-    for (const reroute of reroutes) {
-      reroute.linkIds.add(link.id)
-      if (reroute.floating) reroute.floating = undefined
-      reroute._dragging = undefined
-    }
+    this.finalizeConnectedLink(link)
+    return link
+  }
 
-    // If this is the terminus of a floating link, remove it
-    const lastReroute = reroutes.at(-1)
-    if (lastReroute) {
-      for (const floatingLinkId of lastReroute.floatingLinkIds) {
-        const floatingLink = this.floatingLinks.get(floatingLinkId)
-        if (floatingLink?.parentId === lastReroute.id) {
-          this.removeFloatingLink(floatingLink)
-        }
+  connectSubgraphInputSlot(
+    subgraphInput: SubgraphInput,
+    targetNode: LGraphNode,
+    targetSlotIndex: number,
+    afterRerouteId?: RerouteId
+  ): LLink {
+    const targetInput = targetNode.inputs[targetSlotIndex]
+    const subgraphInputIndex = subgraphInput.parent.slots.indexOf(subgraphInput)
+    const link = new LLink(
+      ++this.state.lastLinkId,
+      targetInput.type,
+      subgraphInput.parent.id,
+      subgraphInputIndex,
+      targetNode.id,
+      targetSlotIndex,
+      afterRerouteId
+    )
+
+    this._links.set(link.id, link)
+    subgraphInput.linkIds.push(link.id)
+    targetInput.link = link.id
+
+    this.finalizeConnectedLink(link)
+    return link
+  }
+
+  connectSubgraphOutputSlot(
+    sourceNode: LGraphNode,
+    sourceSlotIndex: number,
+    subgraphOutput: SubgraphOutput,
+    afterRerouteId?: RerouteId
+  ): LLink {
+    const sourceOutput = sourceNode.outputs[sourceSlotIndex]
+    const subgraphOutputIndex =
+      subgraphOutput.parent.slots.indexOf(subgraphOutput)
+    const link = new LLink(
+      ++this.state.lastLinkId,
+      sourceOutput.type,
+      sourceNode.id,
+      sourceSlotIndex,
+      subgraphOutput.parent.id,
+      subgraphOutputIndex,
+      afterRerouteId
+    )
+
+    this._links.set(link.id, link)
+    subgraphOutput.linkIds[0] = link.id
+    sourceOutput.links ??= []
+    sourceOutput.links.push(link.id)
+
+    this.finalizeConnectedLink(link)
+    return link
+  }
+
+  disconnectSubgraphInputLink(
+    subgraphInput: SubgraphInput,
+    targetNode: LGraphNode,
+    targetSlotIndex: number,
+    link: LLink | undefined
+  ): void {
+    const targetInput = targetNode.inputs[targetSlotIndex]
+    if (targetInput._floatingLinks?.size) {
+      for (const floatingLink of targetInput._floatingLinks) {
+        this.removeFloatingLink(floatingLink)
       }
     }
 
+    targetInput.link = null
+    if (!link) return
+
+    this.disconnectLink(link, 'output')
     this._version++
-    return link
+
+    const index = subgraphInput.linkIds.indexOf(link.id)
+    if (index === -1) {
+      console.warn(
+        'disconnectSubgraphInputLink: link ID not found in subgraph input',
+        link.id
+      )
+      return
+    }
+    subgraphInput.linkIds.splice(index, 1)
+  }
+
+  disconnectSubgraphOutputLink(
+    subgraphOutput: SubgraphOutput,
+    sourceNode: LGraphNode,
+    sourceSlotIndex: number,
+    link: LLink
+  ): void {
+    const sourceOutput = sourceNode.outputs[sourceSlotIndex]
+    this.disconnectLink(link, 'input')
+    this._version++
+
+    if (sourceOutput.links) {
+      sourceOutput.links = sourceOutput.links.filter((id) => id !== link.id)
+    }
+
+    const subgraphLinkIndex = subgraphOutput.linkIds.indexOf(link.id)
+    if (subgraphLinkIndex !== -1) {
+      subgraphOutput.linkIds.splice(subgraphLinkIndex, 1)
+    }
   }
 
   /**

src/lib/litegraph/src/LGraphNode.ts

@@ -3093,7 +3093,20 @@ export class LGraphNode
         ) {
           const targetSlot = graph.outputNode.slots[link_info.target_slot]
           if (targetSlot) {
-            targetSlot.linkIds.length = 0
+            graph.disconnectSubgraphOutputLink(
+              targetSlot,
+              this,
+              slot,
+              link_info
+            )
+            this.onConnectionsChange?.(
+              NodeSlotType.OUTPUT,
+              slot,
+              false,
+              link_info,
+              output
+            )
+            continue
           } else {
             console.error('Missing subgraphOutput slot when disconnecting link')
           }

src/lib/litegraph/src/subgraph/SubgraphIO.test.ts

@@ -1,9 +1,14 @@
 // TODO: Fix these tests after migration
-import { describe, expect, it } from 'vitest'
+import { describe, expect, it, vi } from 'vitest'
 
 import { LGraphNode } from '@/lib/litegraph/src/litegraph'
 import { ToInputFromIoNodeLink } from '@/lib/litegraph/src/canvas/ToInputFromIoNodeLink'
+import {
+  SUBGRAPH_INPUT_ID,
+  SUBGRAPH_OUTPUT_ID
+} from '@/lib/litegraph/src/constants'
 import { LinkDirection } from '@/lib/litegraph/src//types/globalEnums'
+import { NodeSlotType } from '@/lib/litegraph/src/types/globalEnums'
 
 import { subgraphTest } from './__fixtures__/subgraphFixtures'
 import {
@@ -12,6 +17,81 @@ import {
 } from './__fixtures__/subgraphHelpers'
 
 describe('SubgraphIO - Input Slot Dual-Nature Behavior', () => {
+  subgraphTest(
+    'connect callback payload keeps current subgraph-input asymmetry',
+    ({ subgraphWithNode }) => {
+      const { subgraph } = subgraphWithNode
+      const internalNode = new LGraphNode('Internal Target')
+      internalNode.addInput('in', '*')
+      subgraph.add(internalNode)
+
+      const inputCallback = vi.fn()
+      internalNode.onConnectionsChange = inputCallback
+
+      const subgraphInput = subgraph.inputNode.slots[0]
+      const nodeInput = internalNode.inputs[0]
+      const link = subgraphInput.connect(nodeInput, internalNode)
+
+      expect(link).toBeDefined()
+      expect(link?.origin_id).toBe(SUBGRAPH_INPUT_ID)
+      expect(link?.target_id).toBe(internalNode.id)
+      expect(inputCallback).toHaveBeenCalledTimes(1)
+      expect(inputCallback).toHaveBeenLastCalledWith(
+        NodeSlotType.INPUT,
+        0,
+        true,
+        link,
+        nodeInput
+      )
+    }
+  )
+
+  subgraphTest(
+    'disconnect callback payload keeps current subgraph-input asymmetry',
+    ({ subgraphWithNode }) => {
+      const { subgraph } = subgraphWithNode
+      const internalNode = new LGraphNode('Internal Target')
+      internalNode.addInput('in', '*')
+      subgraph.add(internalNode)
+
+      const inputCallback = vi.fn()
+      internalNode.onConnectionsChange = inputCallback
+
+      const subgraphInput = subgraph.inputNode.slots[0]
+      const link = subgraphInput.connect(internalNode.inputs[0], internalNode)
+      if (!link) throw new Error('Expected link')
+
+      new ToInputFromIoNodeLink(
+        subgraph,
+        subgraph.inputNode,
+        subgraphInput,
+        undefined,
+        LinkDirection.CENTER,
+        link
+      ).disconnect()
+
+      expect(inputCallback).toHaveBeenNthCalledWith(
+        1,
+        NodeSlotType.INPUT,
+        0,
+        true,
+        link,
+        internalNode.inputs[0]
+      )
+      expect(inputCallback).toHaveBeenNthCalledWith(
+        2,
+        NodeSlotType.INPUT,
+        0,
+        false,
+        link,
+        subgraphInput
+      )
+      expect(internalNode.inputs[0].link).toBeNull()
+      expect(subgraphInput.linkIds).toEqual([])
+      expect(subgraph.links.get(link.id)).toBeUndefined()
+    }
+  )
+
   subgraphTest(
     'input accepts external connections from parent graph',
     ({ subgraphWithNode }) => {
@@ -128,6 +208,75 @@ describe('SubgraphIO - Input Slot Dual-Nature Behavior', () => {
 })
 
 describe('SubgraphIO - Output Slot Dual-Nature Behavior', () => {
+  subgraphTest(
+    'connect callback payload keeps current subgraph-output asymmetry',
+    ({ subgraphWithNode }) => {
+      const { subgraph } = subgraphWithNode
+
+      const internalNode = new LGraphNode('Internal Source')
+      internalNode.addOutput('out', '*')
+      subgraph.add(internalNode)
+
+      const outputCallback = vi.fn()
+      internalNode.onConnectionsChange = outputCallback
+
+      const subgraphOutput = subgraph.outputNode.slots[0]
+      const nodeOutput = internalNode.outputs[0]
+      const link = subgraphOutput.connect(nodeOutput, internalNode)
+
+      expect(link).toBeDefined()
+      expect(link?.origin_id).toBe(internalNode.id)
+      expect(link?.target_id).toBe(SUBGRAPH_OUTPUT_ID)
+      expect(outputCallback).toHaveBeenLastCalledWith(
+        NodeSlotType.OUTPUT,
+        0,
+        true,
+        link,
+        nodeOutput
+      )
+    }
+  )
+
+  subgraphTest(
+    'disconnect callback payload keeps current subgraph-output asymmetry',
+    ({ subgraphWithNode }) => {
+      const { subgraph } = subgraphWithNode
+
+      const internalNode = new LGraphNode('Internal Source')
+      internalNode.addOutput('out', '*')
+      subgraph.add(internalNode)
+
+      const outputCallback = vi.fn()
+      internalNode.onConnectionsChange = outputCallback
+
+      const subgraphOutput = subgraph.outputNode.slots[0]
+      const link = subgraphOutput.connect(internalNode.outputs[0], internalNode)
+      if (!link) throw new Error('Expected link')
+
+      subgraphOutput.disconnect()
+
+      expect(outputCallback).toHaveBeenNthCalledWith(
+        1,
+        NodeSlotType.OUTPUT,
+        0,
+        true,
+        link,
+        internalNode.outputs[0]
+      )
+      expect(outputCallback).toHaveBeenNthCalledWith(
+        2,
+        NodeSlotType.OUTPUT,
+        0,
+        false,
+        link,
+        subgraphOutput
+      )
+      expect(subgraph.links.get(link.id)).toBeUndefined()
+      expect(subgraphOutput.linkIds).toEqual([])
+      expect(internalNode.outputs[0].links).toEqual([])
+    }
+  )
+
   subgraphTest(
     'output provides connections to parent graph',
     ({ subgraphWithNode }) => {

src/lib/litegraph/src/subgraph/SubgraphInput.ts

@@ -1,5 +1,5 @@
 import type { LGraphNode } from '@/lib/litegraph/src/LGraphNode'
-import { LLink } from '@/lib/litegraph/src/LLink'
+import type { LLink } from '@/lib/litegraph/src/LLink'
 import type { RerouteId } from '@/lib/litegraph/src/Reroute'
 import { CustomEventTarget } from '@/lib/litegraph/src/infrastructure/CustomEventTarget'
 import type { SubgraphInputEventMap } from '@/lib/litegraph/src/infrastructure/SubgraphInputEventMap'
@@ -95,43 +95,13 @@ export class SubgraphInput extends SubgraphSlot {
       })
     }
 
-    const link = new LLink(
-      ++subgraph.state.lastLinkId,
-      slot.type,
-      this.parent.id,
-      this.parent.slots.indexOf(this),
-      node.id,
+    const link = subgraph.connectSubgraphInputSlot(
+      this,
+      node,
       inputIndex,
       afterRerouteId
     )
 
-    // Add to graph links list
-    subgraph._links.set(link.id, link)
-
-    // Set link ID in each slot
-    this.linkIds.push(link.id)
-    slot.link = link.id
-
-    // Reroutes
-    const reroutes = LLink.getReroutes(subgraph, link)
-    for (const reroute of reroutes) {
-      reroute.linkIds.add(link.id)
-      if (reroute.floating) delete reroute.floating
-      reroute._dragging = undefined
-    }
-
-    // If this is the terminus of a floating link, remove it
-    const lastReroute = reroutes.at(-1)
-    if (lastReroute) {
-      for (const linkId of lastReroute.floatingLinkIds) {
-        const link = subgraph.floatingLinks.get(linkId)
-        if (link?.parentId === lastReroute.id) {
-          subgraph.removeFloatingLink(link)
-        }
-      }
-    }
-    subgraph._version++
-
     node.onConnectionsChange?.(NodeSlotType.INPUT, inputIndex, true, link, slot)
 
     subgraph.afterChange()

src/lib/litegraph/src/subgraph/SubgraphInputNode.ts

@@ -170,52 +170,32 @@ export class SubgraphInputNode
   ): void {
     const { subgraph } = this
 
-    // Break floating links
-    if (input._floatingLinks?.size) {
-      for (const link of input._floatingLinks) {
-        subgraph.removeFloatingLink(link)
-      }
-    }
-
-    input.link = null
-    subgraph.setDirtyCanvas(false, true)
-
-    if (!link) return
-
-    const subgraphInputIndex = link.origin_slot
-    link.disconnect(subgraph, 'output')
-    subgraph._version++
-
-    const subgraphInput = this.slots.at(subgraphInputIndex)
+    const subgraphInput = link ? this.slots.at(link.origin_slot) : undefined
     if (!subgraphInput) {
       console.warn(
         'disconnectNodeInput: subgraphInput not found',
         this,
-        subgraphInputIndex
+        link?.origin_slot
       )
       return
     }
 
-    // search in the inputs list for this link
-    const index = subgraphInput.linkIds.indexOf(link.id)
-    if (index !== -1) {
-      subgraphInput.linkIds.splice(index, 1)
-    } else {
-      console.warn(
-        'disconnectNodeInput: link ID not found in subgraphInput linkIds',
-        link.id
-      )
-    }
     const slotIndex = node.inputs.findIndex((inp) => inp === input)
-    if (slotIndex !== -1) {
-      node.onConnectionsChange?.(
-        NodeSlotType.INPUT,
-        slotIndex,
-        false,
-        link,
-        subgraphInput
-      )
+    if (slotIndex === -1) {
+      console.warn('disconnectNodeInput: target input slot not found', this)
+      return
     }
+
+    subgraph.disconnectSubgraphInputLink(subgraphInput, node, slotIndex, link)
+    subgraph.setDirtyCanvas(false, true)
+
+    node.onConnectionsChange?.(
+      NodeSlotType.INPUT,
+      slotIndex,
+      false,
+      link,
+      subgraphInput
+    )
   }
 
   override drawProtected(

src/lib/litegraph/src/subgraph/SubgraphOutput.ts

@@ -1,7 +1,5 @@
-import { pull } from 'es-toolkit/compat'
-
 import type { LGraphNode } from '@/lib/litegraph/src/LGraphNode'
-import { LLink } from '@/lib/litegraph/src/LLink'
+import type { LLink } from '@/lib/litegraph/src/LLink'
 import type { RerouteId } from '@/lib/litegraph/src/Reroute'
 import type {
   INodeInputSlot,
@@ -57,50 +55,24 @@ export class SubgraphOutput extends SubgraphSlot {
     if (existingLink != null) {
       subgraph.beforeChange()
 
-      existingLink.disconnect(subgraph, 'input')
-      const resolved = existingLink.resolve(subgraph)
-      const links = resolved.output?.links
-      if (links) pull(links, existingLink.id)
+      const { outputNode } = existingLink.resolve(subgraph)
+      if (!outputNode) throw new Error('Expected output node for existing link')
+
+      subgraph.disconnectSubgraphOutputLink(
+        this,
+        outputNode,
+        existingLink.origin_slot,
+        existingLink
+      )
     }
 
-    const link = new LLink(
-      ++subgraph.state.lastLinkId,
-      slot.type,
-      node.id,
+    const link = subgraph.connectSubgraphOutputSlot(
+      node,
       outputIndex,
-      this.parent.id,
-      this.parent.slots.indexOf(this),
+      this,
       afterRerouteId
     )
 
-    // Add to graph links list
-    subgraph._links.set(link.id, link)
-
-    // Set link ID in each slot
-    this.linkIds[0] = link.id
-    slot.links ??= []
-    slot.links.push(link.id)
-
-    // Reroutes
-    const reroutes = LLink.getReroutes(subgraph, link)
-    for (const reroute of reroutes) {
-      reroute.linkIds.add(link.id)
-      if (reroute.floating) delete reroute.floating
-      reroute._dragging = undefined
-    }
-
-    // If this is the terminus of a floating link, remove it
-    const lastReroute = reroutes.at(-1)
-    if (lastReroute) {
-      for (const linkId of lastReroute.floatingLinkIds) {
-        const link = subgraph.floatingLinks.get(linkId)
-        if (link?.parentId === lastReroute.id) {
-          subgraph.removeFloatingLink(link)
-        }
-      }
-    }
-    subgraph._version++
-
     node.onConnectionsChange?.(
       NodeSlotType.OUTPUT,
       outputIndex,
@@ -156,13 +128,20 @@ export class SubgraphOutput extends SubgraphSlot {
   override disconnect() {
     const { subgraph } = this.parent
     //should never have more than one connection
-    for (const linkId of this.linkIds) {
+    for (const linkId of [...this.linkIds]) {
       const link = subgraph.links[linkId]
       if (!link) continue
-      subgraph.removeLink(linkId)
-      const { output, outputNode } = link.resolve(subgraph)
-      if (output)
-        output.links = output.links?.filter((id) => id !== linkId) ?? null
+
+      const { outputNode } = link.resolve(subgraph)
+      if (!outputNode) continue
+
+      subgraph.disconnectSubgraphOutputLink(
+        this,
+        outputNode,
+        link.origin_slot,
+        link
+      )
+
       outputNode?.onConnectionsChange?.(
         NodeSlotType.OUTPUT,
         link.origin_slot,
@@ -171,6 +150,5 @@ export class SubgraphOutput extends SubgraphSlot {
         this
       )
     }
-    this.linkIds.length = 0
   }
 }