Fix links being lost after manage group node (#916)

* Fix links being lost after manage group node

* Change to use groupnodebuilder

* Make test more reliable

browser_tests/ComfyPage.ts

@@ -4,6 +4,8 @@ import dotenv from 'dotenv'
 dotenv.config()
 import * as fs from 'fs'
 import { NodeBadgeMode } from '../src/types/nodeSource'
+import { NodeId } from '../src/types/comfyWorkflow'
+import { ManageGroupNode } from './helpers/manageGroupNode'
 
 interface Position {
   x: number
@@ -707,6 +709,177 @@ export class ComfyPage {
     await this.page.getByText('Convert to Group Node').click()
     await this.nextFrame()
   }
+  async convertOffsetToCanvas(pos: [number, number]) {
+    return this.page.evaluate((pos) => {
+      return window['app'].canvas.ds.convertOffsetToCanvas(pos)
+    }, pos)
+  }
+  async getNodeRefById(id: NodeId) {
+    return new NodeReference(id, this)
+  }
+  async getNodeRefsByType(type: string): Promise<NodeReference[]> {
+    return (
+      await this.page.evaluate((type) => {
+        return window['app'].graph._nodes
+          .filter((n) => n.type === type)
+          .map((n) => n.id)
+      }, type)
+    ).map((id: NodeId) => this.getNodeRefById(id))
+  }
+}
+class NodeSlotReference {
+  constructor(
+    readonly type: 'input' | 'output',
+    readonly index: number,
+    readonly node: NodeReference
+  ) {}
+  async getPosition() {
+    const pos: [number, number] = await this.node.comfyPage.page.evaluate(
+      ([type, id, index]) => {
+        const node = window['app'].graph.getNodeById(id)
+        if (!node) throw new Error(`Node ${id} not found.`)
+        return window['app'].canvas.ds.convertOffsetToCanvas(
+          node.getConnectionPos(type === 'input', index)
+        )
+      },
+      [this.type, this.node.id, this.index] as const
+    )
+    return {
+      x: pos[0],
+      y: pos[1]
+    }
+  }
+  async getLinkCount() {
+    return await this.node.comfyPage.page.evaluate(
+      ([type, id, index]) => {
+        const node = window['app'].graph.getNodeById(id)
+        if (!node) throw new Error(`Node ${id} not found.`)
+        if (type === 'input') {
+          return node.inputs[index].link == null ? 0 : 1
+        }
+        return node.outputs[index].links?.length ?? 0
+      },
+      [this.type, this.node.id, this.index] as const
+    )
+  }
+  async removeLinks() {
+    await this.node.comfyPage.page.evaluate(
+      ([type, id, index]) => {
+        const node = window['app'].graph.getNodeById(id)
+        if (!node) throw new Error(`Node ${id} not found.`)
+        if (type === 'input') {
+          node.disconnectInput(index)
+        } else {
+          node.disconnectOutput(index)
+        }
+      },
+      [this.type, this.node.id, this.index] as const
+    )
+  }
+}
+class NodeReference {
+  constructor(
+    readonly id: NodeId,
+    readonly comfyPage: ComfyPage
+  ) {}
+  async exists(): Promise<boolean> {
+    return await this.comfyPage.page.evaluate((id) => {
+      const node = window['app'].graph.getNodeById(id)
+      return !!node
+    }, this.id)
+  }
+  getType(): Promise<string> {
+    return this.getProperty('type')
+  }
+  async getPosition(): Promise<Position> {
+    const pos = await this.comfyPage.convertOffsetToCanvas(
+      await this.getProperty<[number, number]>('pos')
+    )
+    return {
+      x: pos[0],
+      y: pos[1]
+    }
+  }
+  async getSize(): Promise<Size> {
+    const size = await this.getProperty<[number, number]>('size')
+    return {
+      width: size[0],
+      height: size[1]
+    }
+  }
+  async getProperty<T>(prop: string): Promise<T> {
+    return await this.comfyPage.page.evaluate(
+      ([id, prop]) => {
+        const node = window['app'].graph.getNodeById(id)
+        if (!node) throw new Error('Node not found')
+        return node[prop]
+      },
+      [this.id, prop] as const
+    )
+  }
+  async getOutput(index: number) {
+    return new NodeSlotReference('output', index, this)
+  }
+  async getInput(index: number) {
+    return new NodeSlotReference('input', index, this)
+  }
+  async click(position: 'title', options?: Parameters<Page['click']>[1]) {
+    const nodePos = await this.getPosition()
+    const nodeSize = await this.getSize()
+    let clickPos: Position
+    switch (position) {
+      case 'title':
+        clickPos = { x: nodePos.x + nodeSize.width / 2, y: nodePos.y + 15 }
+        break
+      default:
+        throw new Error(`Invalid click position ${position}`)
+    }
+    await this.comfyPage.canvas.click({
+      ...options,
+      position: clickPos
+    })
+    await this.comfyPage.nextFrame()
+  }
+  async connectOutput(
+    originSlotIndex: number,
+    targetNode: NodeReference,
+    targetSlotIndex: number
+  ) {
+    const originSlot = await this.getOutput(originSlotIndex)
+    const targetSlot = await targetNode.getInput(targetSlotIndex)
+    await this.comfyPage.dragAndDrop(
+      await originSlot.getPosition(),
+      await targetSlot.getPosition()
+    )
+    return originSlot
+  }
+  async clickContextMenuOption(optionText: string) {
+    await this.click('title', { button: 'right' })
+    const ctx = this.comfyPage.page.locator('.litecontextmenu')
+    await ctx.getByText(optionText).click()
+  }
+  async convertToGroupNode(groupNodeName: string = 'GroupNode') {
+    this.comfyPage.page.once('dialog', async (dialog) => {
+      await dialog.accept(groupNodeName)
+    })
+    await this.clickContextMenuOption('Convert to Group Node')
+    await this.comfyPage.nextFrame()
+    const nodes = await this.comfyPage.getNodeRefsByType(
+      `workflow/${groupNodeName}`
+    )
+    if (nodes.length !== 1) {
+      throw new Error(`Did not find single group node (found=${nodes.length})`)
+    }
+    return nodes[0]
+  }
+  async manageGroupNode() {
+    await this.clickContextMenuOption('Manage Group Node')
+    await this.comfyPage.nextFrame()
+    return new ManageGroupNode(
+      this.comfyPage.page,
+      this.comfyPage.page.locator('.comfy-group-manage')
+    )
+  }
 }
 
 export const comfyPageFixture = base.extend<{ comfyPage: ComfyPage }>({

browser_tests/groupNode.spec.ts

@@ -53,4 +53,40 @@ test.describe('Group Node', () => {
     await comfyPage.page.waitForTimeout(tooltipTimeout + 16)
     await expect(comfyPage.page.locator('.node-tooltip')).toBeVisible()
   })
+
+  test('Reconnects inputs after configuration changed via manage dialog save', async ({
+    comfyPage
+  }) => {
+    const expectSingleNode = async (type: string) => {
+      const nodes = await comfyPage.getNodeRefsByType(type)
+      expect(nodes).toHaveLength(1)
+      return nodes[0]
+    }
+    const latent = await expectSingleNode('EmptyLatentImage')
+    const sampler = await expectSingleNode('KSampler')
+    // Remove existing link
+    const samplerInput = await sampler.getInput(0)
+    await samplerInput.removeLinks()
+    // Group latent + sampler
+    await latent.click('title', {
+      modifiers: ['Shift']
+    })
+    await sampler.click('title', {
+      modifiers: ['Shift']
+    })
+    const groupNode = await sampler.convertToGroupNode()
+    // Connect node to group
+    const ckpt = await expectSingleNode('CheckpointLoaderSimple')
+    const input = await ckpt.connectOutput(0, groupNode, 0)
+    expect(await input.getLinkCount()).toBe(1)
+    // Modify the group node via manage dialog
+    const manage = await groupNode.manageGroupNode()
+    await manage.selectNode('KSampler')
+    await manage.changeTab('Inputs')
+    await manage.setLabel('model', 'test')
+    await manage.save()
+    await manage.close()
+    // Ensure the link is still present
+    expect(await input.getLinkCount()).toBe(1)
+  })
 })

browser_tests/helpers/manageGroupNode.ts

@@ -0,0 +1,37 @@
+import { Locator, Page } from '@playwright/test'
+export class ManageGroupNode {
+  footer: Locator
+
+  constructor(
+    readonly page: Page,
+    readonly root: Locator
+  ) {
+    this.footer = root.locator('footer')
+  }
+
+  async setLabel(name: string, label: string) {
+    const active = this.root.locator('.comfy-group-manage-node-page.active')
+    const input = active.getByPlaceholder(name)
+    await input.fill(label)
+  }
+
+  async save() {
+    await this.footer.getByText('Save').click()
+  }
+
+  async close() {
+    await this.footer.getByText('Close').click()
+  }
+
+  async selectNode(name: string) {
+    const list = this.root.locator('.comfy-group-manage-list-items')
+    const item = list.getByText(name)
+    await item.click()
+  }
+
+  async changeTab(name: 'Inputs' | 'Widgets' | 'Outputs') {
+    const header = this.root.locator('.comfy-group-manage-node header')
+    const tab = header.getByText(name)
+    await tab.click()
+  }
+}

src/extensions/core/groupNode.ts

@@ -828,6 +828,9 @@ export class GroupNodeHandler {
       ]
 
       // Remove all converted nodes and relink them
+      const builder = new GroupNodeBuilder(nodes)
+      const nodeData = builder.getNodeData()
+      groupNode[GROUP].groupData.nodeData.links = nodeData.links
       groupNode[GROUP].replaceNodes(nodes)
       return groupNode
     }