feat: extract shared subgraph boundary remap adapter

Amp-Thread-ID: https://ampcode.com/threads/T-019c98bb-ebfb-77fc-b105-1d04e023db7a
Co-authored-by: Amp <amp@ampcode.com>

src/lib/litegraph/src/LGraph.test.ts

@@ -7,7 +7,8 @@ import {
   LGraph,
   LGraphNode,
   LiteGraph,
-  LLink
+  LLink,
+  SubgraphNode
 } from '@/lib/litegraph/src/litegraph'
 import { NodeSlotType } from '@/lib/litegraph/src/types/globalEnums'
 import type { UUID } from '@/lib/litegraph/src/utils/uuid'
@@ -786,6 +787,42 @@ describe('ensureGlobalIdUniqueness', () => {
 })
 
 describe('Subgraph Unpacking', () => {
+  function installSubgraphNodeRegistration(rootGraph: LGraph): () => void {
+    const listener = (event: CustomEvent<{ subgraph: Subgraph }>): void => {
+      const { subgraph } = event.detail
+
+      class RuntimeSubgraphNode extends SubgraphNode {
+        constructor(title?: string) {
+          super(rootGraph, subgraph, {
+            id: ++rootGraph.last_node_id,
+            type: subgraph.id,
+            title,
+            pos: [0, 0],
+            size: [140, 80],
+            inputs: [],
+            outputs: [],
+            properties: {},
+            flags: {},
+            mode: 0,
+            order: 0
+          })
+        }
+      }
+
+      LiteGraph.registerNodeType(subgraph.id, RuntimeSubgraphNode)
+    }
+
+    rootGraph.events.addEventListener('subgraph-created', listener)
+    return () =>
+      rootGraph.events.removeEventListener('subgraph-created', listener)
+  }
+
+  function getRequiredNodeByTitle(graph: LGraph, title: string): LGraphNode {
+    const node = graph.nodes.find((candidate) => candidate.title === title)
+    if (!node) throw new Error(`Expected node titled ${title}`)
+    return node
+  }
+
   class TestNode extends LGraphNode {
     constructor(title?: string) {
       super(title ?? 'TestNode')
@@ -890,4 +927,115 @@ describe('Subgraph Unpacking', () => {
     expect(unpackedTarget.inputs[0].link).not.toBeNull()
     expect(unpackedTarget.inputs[1].link).toBeNull()
   })
+
+  it('preserves boundary input reroute parent remap across convert and unpack', () => {
+    setActivePinia(createTestingPinia({ stubActions: false }))
+    registerTestNodes()
+    const rootGraph = new LGraph()
+    const cleanupRegistration = installSubgraphNodeRegistration(rootGraph)
+    try {
+      const externalSource = LiteGraph.createNode(
+        'test/TestNode',
+        'external-source'
+      )
+      const boundaryTarget = LiteGraph.createNode(
+        'test/TestNode',
+        'boundary-target'
+      )
+      if (!externalSource || !boundaryTarget)
+        throw new Error('Expected test nodes')
+      rootGraph.add(externalSource)
+      rootGraph.add(boundaryTarget)
+
+      const boundaryLink = externalSource.connect(0, boundaryTarget, 0)
+      if (!boundaryLink) throw new Error('Expected boundary link')
+
+      const reroute = rootGraph.createReroute([120, 40], boundaryLink)
+      expect(boundaryLink.parentId).toBe(reroute.id)
+
+      const { node: subgraphNode } = rootGraph.convertToSubgraph(
+        new Set([boundaryTarget])
+      )
+      const convertedBoundaryLinkId = subgraphNode.inputs[0].link
+      if (convertedBoundaryLinkId == null)
+        throw new Error('Expected converted boundary input link')
+
+      const convertedBoundaryLink = rootGraph.getLink(convertedBoundaryLinkId)
+      if (!convertedBoundaryLink)
+        throw new Error('Expected converted boundary input link instance')
+      expect(convertedBoundaryLink.parentId).toBe(reroute.id)
+
+      rootGraph.unpackSubgraph(subgraphNode)
+
+      const unpackedTarget = getRequiredNodeByTitle(
+        rootGraph,
+        'boundary-target'
+      )
+      const unpackedLink = rootGraph.getLink(unpackedTarget.inputs[0].link)
+      if (!unpackedLink)
+        throw new Error('Expected unpacked boundary input link')
+
+      expect(unpackedLink.origin_id).toBe(externalSource.id)
+      expect(unpackedLink.target_id).toBe(unpackedTarget.id)
+      expect(unpackedLink.parentId).toBe(reroute.id)
+    } finally {
+      cleanupRegistration()
+    }
+  })
+
+  it('preserves boundary output reroute parent remap across convert and unpack', () => {
+    setActivePinia(createTestingPinia({ stubActions: false }))
+    registerTestNodes()
+    const rootGraph = new LGraph()
+    const cleanupRegistration = installSubgraphNodeRegistration(rootGraph)
+    try {
+      const boundarySource = LiteGraph.createNode(
+        'test/TestNode',
+        'boundary-source'
+      )
+      const externalTarget = LiteGraph.createNode(
+        'test/TestNode',
+        'external-target'
+      )
+      if (!boundarySource || !externalTarget)
+        throw new Error('Expected test nodes')
+      rootGraph.add(boundarySource)
+      rootGraph.add(externalTarget)
+
+      const boundaryLink = boundarySource.connect(0, externalTarget, 0)
+      if (!boundaryLink) throw new Error('Expected boundary link')
+
+      const reroute = rootGraph.createReroute([180, 80], boundaryLink)
+      expect(boundaryLink.parentId).toBe(reroute.id)
+
+      const { node: subgraphNode } = rootGraph.convertToSubgraph(
+        new Set([boundarySource])
+      )
+      const convertedBoundaryLinkId = subgraphNode.outputs[0].links?.[0]
+      if (convertedBoundaryLinkId == null)
+        throw new Error('Expected converted boundary output link')
+
+      const convertedBoundaryLink = rootGraph.getLink(convertedBoundaryLinkId)
+      if (!convertedBoundaryLink)
+        throw new Error('Expected converted boundary output link instance')
+      expect(convertedBoundaryLink.parentId).toBe(reroute.id)
+
+      rootGraph.unpackSubgraph(subgraphNode)
+
+      const unpackedSource = getRequiredNodeByTitle(
+        rootGraph,
+        'boundary-source'
+      )
+      const unpackedLinkId = unpackedSource.outputs[0].links?.[0]
+      const unpackedLink = rootGraph.getLink(unpackedLinkId)
+      if (!unpackedLink)
+        throw new Error('Expected unpacked boundary output link')
+
+      expect(unpackedLink.origin_id).toBe(unpackedSource.id)
+      expect(unpackedLink.target_id).toBe(externalTarget.id)
+      expect(unpackedLink.parentId).toBe(reroute.id)
+    } finally {
+      cleanupRegistration()
+    }
+  })
 })

src/lib/litegraph/src/LGraph.ts

@@ -59,6 +59,7 @@ import {
   mapSubgraphInputsAndLinks,
   mapSubgraphOutputsAndLinks,
   multiClone,
+  subgraphBoundaryAdapter,
   splitPositionables
 } from './subgraph/subgraphUtils'
 import { Alignment, LGraphEventMode, NodeSlotType } from './types/globalEnums'
@@ -1991,9 +1992,12 @@ export class LGraph
 
       // Special handling: Subgraph input node
       i++
-      if (link.origin_id === SUBGRAPH_INPUT_ID) {
-        link.target_id = subgraphNode.id
-        link.target_slot = i - 1
+      if (subgraphBoundaryAdapter.isInputBoundary(link)) {
+        subgraphBoundaryAdapter.remapInputBoundaryForConvert(
+          link,
+          subgraphNode.id,
+          i - 1
+        )
         if (subgraphInput instanceof SubgraphInput) {
           subgraphInput.connect(
             subgraphNode.findInputSlotByType(link.type, true, true),
@@ -2032,9 +2036,12 @@ export class LGraph
       for (const connection of connections) {
         const { input, inputNode, link, subgraphOutput } = connection
         // Special handling: Subgraph output node
-        if (link.target_id === SUBGRAPH_OUTPUT_ID) {
-          link.origin_id = subgraphNode.id
-          link.origin_slot = i - 1
+        if (subgraphBoundaryAdapter.isOutputBoundary(link)) {
+          subgraphBoundaryAdapter.remapOutputBoundaryForConvert(
+            link,
+            subgraphNode.id,
+            i - 1
+          )
           this.links.set(link.id, link)
           if (subgraphOutput instanceof SubgraphOutput) {
             subgraphOutput.connect(
@@ -2203,16 +2210,20 @@ export class LGraph
     }[] = []
     for (const [, link] of subgraphNode.subgraph._links) {
       let externalParentId: RerouteId | undefined
-      if (link.origin_id === SUBGRAPH_INPUT_ID) {
-        const outerLinkId = subgraphNode.inputs[link.origin_slot].link
-        if (!outerLinkId) {
+      if (subgraphBoundaryAdapter.isInputBoundary(link)) {
+        const endpoint = subgraphBoundaryAdapter.remapInputBoundaryForUnpack(
+          link,
+          subgraphNode,
+          this.links
+        )
+        if (!endpoint) {
           console.error('Missing Link ID when unpacking')
           continue
         }
-        const outerLink = this.links[outerLinkId]
-        link.origin_id = outerLink.origin_id
-        link.origin_slot = outerLink.origin_slot
-        externalParentId = outerLink.parentId
+
+        link.origin_id = endpoint.originId
+        link.origin_slot = endpoint.originSlot
+        externalParentId = endpoint.externalParentId
       } else {
         const origin_id = nodeIdMap.get(link.origin_id)
         if (!origin_id) {
@@ -2221,22 +2232,37 @@ export class LGraph
         }
         link.origin_id = origin_id
       }
-      if (link.target_id === SUBGRAPH_OUTPUT_ID) {
-        for (const linkId of subgraphNode.outputs[link.target_slot].links ??
-          []) {
-          const sublink = this.links[linkId]
+      if (subgraphBoundaryAdapter.isOutputBoundary(link)) {
+        const outputEndpoints =
+          subgraphBoundaryAdapter.resolveOutputBoundaryForUnpack(
+            link,
+            subgraphNode,
+            this.links
+          )
+        if (outputEndpoints.length === 0) {
+          console.error('Missing Link ID when unpacking')
+          continue
+        }
+
+        for (const endpoint of outputEndpoints) {
           newLinks.push({
             oid: link.origin_id,
             oslot: link.origin_slot,
-            tid: sublink.target_id,
-            tslot: sublink.target_slot,
+            tid: endpoint.targetId,
+            tslot: endpoint.targetSlot,
             id: link.id,
             iparent: link.parentId,
-            eparent: sublink.parentId,
+            eparent: endpoint.externalParentId,
             externalFirst: true
           })
-          sublink.parentId = undefined
         }
+
+        for (const linkId of subgraphNode.outputs[link.target_slot].links ??
+          []) {
+          const sublink = this.links.get(linkId)
+          if (sublink) sublink.parentId = undefined
+        }
+
         continue
       } else {
         const target_id = nodeIdMap.get(link.target_id)

src/lib/litegraph/src/subgraph/subgraphUtils.ts

@@ -1,8 +1,9 @@
 import type { LGraph } from '@/lib/litegraph/src/LGraph'
 import { LGraphGroup } from '@/lib/litegraph/src/LGraphGroup'
 import { LGraphNode } from '@/lib/litegraph/src/LGraphNode'
+import type { NodeId } from '@/lib/litegraph/src/LGraphNode'
 import { LLink } from '@/lib/litegraph/src/LLink'
-import type { ResolvedConnection } from '@/lib/litegraph/src/LLink'
+import type { LinkId, ResolvedConnection } from '@/lib/litegraph/src/LLink'
 import { Reroute } from '@/lib/litegraph/src/Reroute'
 import type { RerouteId } from '@/lib/litegraph/src/Reroute'
 import {
@@ -90,6 +91,90 @@ interface BoundaryLinks {
   boundaryOutputLinks: LLink[]
 }
 
+interface SubgraphBoundaryNodeView {
+  id: NodeId
+  inputs: Array<{ link?: LinkId | null }>
+  outputs: Array<{ links?: LinkId[] | null }>
+}
+
+interface SubgraphBoundaryOutputEndpoint {
+  targetId: NodeId
+  targetSlot: number
+  externalParentId: RerouteId | undefined
+}
+
+interface SubgraphBoundaryInputEndpoint {
+  originId: NodeId
+  originSlot: number
+  externalParentId: RerouteId | undefined
+}
+
+export const subgraphBoundaryAdapter = {
+  isInputBoundary(link: LLink): boolean {
+    return link.origin_id === SUBGRAPH_INPUT_ID
+  },
+
+  isOutputBoundary(link: LLink): boolean {
+    return link.target_id === SUBGRAPH_OUTPUT_ID
+  },
+
+  remapInputBoundaryForConvert(
+    link: LLink,
+    subgraphNodeId: NodeId,
+    subgraphInputSlot: number
+  ): void {
+    link.target_id = subgraphNodeId
+    link.target_slot = subgraphInputSlot
+  },
+
+  remapOutputBoundaryForConvert(
+    link: LLink,
+    subgraphNodeId: NodeId,
+    subgraphOutputSlot: number
+  ): void {
+    link.origin_id = subgraphNodeId
+    link.origin_slot = subgraphOutputSlot
+  },
+
+  remapInputBoundaryForUnpack(
+    link: LLink,
+    subgraphNode: SubgraphBoundaryNodeView,
+    links: Map<LinkId, LLink>
+  ): SubgraphBoundaryInputEndpoint | undefined {
+    const outerLinkId = subgraphNode.inputs[link.origin_slot]?.link
+    if (outerLinkId == null) return
+
+    const outerLink = links.get(outerLinkId)
+    if (!outerLink) return
+
+    return {
+      originId: outerLink.origin_id,
+      originSlot: outerLink.origin_slot,
+      externalParentId: outerLink.parentId
+    }
+  },
+
+  resolveOutputBoundaryForUnpack(
+    link: LLink,
+    subgraphNode: SubgraphBoundaryNodeView,
+    links: Map<LinkId, LLink>
+  ): SubgraphBoundaryOutputEndpoint[] {
+    const results: SubgraphBoundaryOutputEndpoint[] = []
+    for (const linkId of subgraphNode.outputs[link.target_slot]?.links ?? []) {
+      const outerLink = links.get(linkId)
+      if (!outerLink) continue
+
+      results.push({
+        targetId: outerLink.target_id,
+        targetSlot: outerLink.target_slot,
+        externalParentId: outerLink.parentId
+      })
+    }
+
+    return results
+  }
+}
+
 export function getBoundaryLinks(
   graph: LGraph,
   items: Set<Positionable>