From 5c6eecd660275427d492742a6213aee70b5eadd9 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Mon, 2 Dec 2024 00:52:08 +0000
Subject: [PATCH] Fix group node manage opening to wrong node type (#1754)

Remove dialog from DOM when closed
Add test
---
 browser_tests/groupNode.spec.ts          | 30 ++++++++++++++++++++++++
 browser_tests/helpers/manageGroupNode.ts |  7 ++++++
 src/extensions/core/groupNode.ts         |  6 ++---
 src/extensions/core/groupNodeManage.ts   |  1 +
 4 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/browser_tests/groupNode.spec.ts b/browser_tests/groupNode.spec.ts
index bb3505ee65..831c1cb437 100644
--- a/browser_tests/groupNode.spec.ts
+++ b/browser_tests/groupNode.spec.ts
@@ -103,6 +103,36 @@ test.describe('Group Node', () => {
     await expect(comfyPage.page.locator('.node-tooltip')).toBeVisible()
   })
 
+  test('Manage group opens with the correct group selected', async ({
+    comfyPage
+  }) => {
+    const makeGroup = async (name, type1, type2) => {
+      const node1 = (await comfyPage.getNodeRefsByType(type1))[0]
+      const node2 = (await comfyPage.getNodeRefsByType(type2))[0]
+      await node1.click('title')
+      await node2.click('title', {
+        modifiers: ['Shift']
+      })
+      return await node2.convertToGroupNode(name)
+    }
+
+    const group1 = await makeGroup(
+      'g1',
+      'CLIPTextEncode',
+      'CheckpointLoaderSimple'
+    )
+    const group2 = await makeGroup('g2', 'EmptyLatentImage', 'KSampler')
+
+    const manage1 = await group1.manageGroupNode()
+    await comfyPage.nextFrame()
+    expect(await manage1.getSelectedNodeType()).toBe('g1')
+    await manage1.close()
+    await expect(manage1.root).not.toBeVisible()
+
+    const manage2 = await group2.manageGroupNode()
+    expect(await manage2.getSelectedNodeType()).toBe('g2')
+  })
+
   test('Reconnects inputs after configuration changed via manage dialog save', async ({
     comfyPage
   }) => {
diff --git a/browser_tests/helpers/manageGroupNode.ts b/browser_tests/helpers/manageGroupNode.ts
index 27bc98346c..beebabd695 100644
--- a/browser_tests/helpers/manageGroupNode.ts
+++ b/browser_tests/helpers/manageGroupNode.ts
@@ -1,12 +1,14 @@
 import { Locator, Page } from '@playwright/test'
 export class ManageGroupNode {
   footer: Locator
+  header: Locator
 
   constructor(
     readonly page: Page,
     readonly root: Locator
   ) {
     this.footer = root.locator('footer')
+    this.header = root.locator('header')
   }
 
   async setLabel(name: string, label: string) {
@@ -23,6 +25,11 @@ export class ManageGroupNode {
     await this.footer.getByText('Close').click()
   }
 
+  async getSelectedNodeType() {
+    const select = this.header.locator('select').first()
+    return await select.inputValue()
+  }
+
   async selectNode(name: string) {
     const list = this.root.locator('.comfy-group-manage-list-items')
     const item = list.getByText(name)
diff --git a/src/extensions/core/groupNode.ts b/src/extensions/core/groupNode.ts
index 5905bf00ea..8aeec594af 100644
--- a/src/extensions/core/groupNode.ts
+++ b/src/extensions/core/groupNode.ts
@@ -1001,7 +1001,7 @@ export class GroupNodeHandler {
         },
         {
           content: 'Manage Group Node',
-          callback: manageGroupNodes
+          callback: () => manageGroupNodes(this.type)
         }
       )
     }
@@ -1488,8 +1488,8 @@ function ungroupSelectedGroupNodes() {
   }
 }
 
-function manageGroupNodes() {
-  new ManageGroupDialog(app).show()
+function manageGroupNodes(type?: string) {
+  new ManageGroupDialog(app).show(type)
 }
 
 const id = 'Comfy.GroupNode'
diff --git a/src/extensions/core/groupNodeManage.ts b/src/extensions/core/groupNodeManage.ts
index 61dfa6d2e2..d845bd4e4d 100644
--- a/src/extensions/core/groupNodeManage.ts
+++ b/src/extensions/core/groupNodeManage.ts
@@ -519,6 +519,7 @@ export class ManageGroupDialog extends ComfyDialog<HTMLDialogElement> {
 
     this.element.addEventListener('close', () => {
       this.draggable?.dispose()
+      this.element.remove()
     })
   }
 }