Auto link node on creation (#110)

* Auto link node on creation

* Handle corner case

* Add some browser tests

* Add auto link test

* Force enable

* Confirm setting before running test

* Update test expectations [skip ci]

---------

Co-authored-by: github-actions <[email protected]>

Author: huchenlei

browser_tests/ComfyPage.ts

@@ -1,5 +1,5 @@
-import type { Page, Locator } from '@playwright/test';
-import { test as base } from '@playwright/test';
+import type { Page, Locator } from "@playwright/test";
+import { test as base } from "@playwright/test";
 import dotenv from "dotenv";
 dotenv.config();
 
@@ -8,6 +8,27 @@ interface Position {
   y: number;
 }
 
+class ComfyNodeSearchBox {
+  public readonly input: Locator;
+  public readonly dropdown: Locator;
+
+  constructor(public readonly page: Page) {
+    this.input = page.locator(
+      '.comfy-vue-node-search-container input[type="text"]'
+    );
+    this.dropdown = page.locator(
+      ".comfy-vue-node-search-container .p-autocomplete-list"
+    );
+  }
+
+  async fillAndSelectFirstNode(nodeName: string) {
+    await this.input.waitFor({ state: "visible" });
+    await this.input.fill(nodeName);
+    await this.dropdown.waitFor({ state: "visible" });
+    await this.dropdown.locator("li").nth(0).click();
+  }
+}
+
 export class ComfyPage {
   public readonly url: string;
   // All canvas position operations are based on default view of canvas.
@@ -17,13 +38,15 @@ export class ComfyPage {
   // Buttons
   public readonly resetViewButton: Locator;
 
-  constructor(
-    public readonly page: Page,
-  ) {
-    this.url = process.env.PLAYWRIGHT_TEST_URL || 'http://localhost:8188';
-    this.canvas = page.locator('#graph-canvas');
-    this.widgetTextBox = page.getByPlaceholder('text').nth(1);
-    this.resetViewButton = page.getByRole('button', { name: 'Reset View' });
+  // Search box
+  public readonly searchBox: ComfyNodeSearchBox;
+
+  constructor(public readonly page: Page) {
+    this.url = process.env.PLAYWRIGHT_TEST_URL || "http://localhost:8188";
+    this.canvas = page.locator("#graph-canvas");
+    this.widgetTextBox = page.getByPlaceholder("text").nth(1);
+    this.resetViewButton = page.getByRole("button", { name: "Reset View" });
+    this.searchBox = new ComfyNodeSearchBox(page);
   }
 
   async goto() {
@@ -47,8 +70,8 @@ export class ComfyPage {
     await this.canvas.click({
       position: {
         x: 618,
-        y: 191
-      }
+        y: 191,
+      },
     });
     await this.nextFrame();
   }
@@ -57,8 +80,8 @@ export class ComfyPage {
     await this.canvas.click({
       position: {
         x: 622,
-        y: 400
-      }
+        y: 400,
+      },
     });
     await this.nextFrame();
   }
@@ -67,8 +90,8 @@ export class ComfyPage {
     await this.canvas.click({
       position: {
         x: 35,
-        y: 31
-      }
+        y: 31,
+      },
     });
     await this.nextFrame();
   }
@@ -82,10 +105,7 @@ export class ComfyPage {
   }
 
   async dragNode2() {
-    await this.dragAndDrop(
-      { x: 622, y: 400 },
-      { x: 622, y: 300 },
-    );
+    await this.dragAndDrop({ x: 622, y: 400 }, { x: 622, y: 300 });
     await this.nextFrame();
   }
 
@@ -113,15 +133,15 @@ export class ComfyPage {
   async adjustWidgetValue() {
     // Adjust Empty Latent Image's width input.
     const page = this.page;
-    await page.locator('#graph-canvas').click({
+    await page.locator("#graph-canvas").click({
       position: {
         x: 724,
-        y: 645
-      }
+        y: 645,
+      },
     });
     await page.locator('input[type="text"]').click();
-    await page.locator('input[type="text"]').fill('128');
-    await page.locator('input[type="text"]').press('Enter');
+    await page.locator('input[type="text"]').fill("128");
+    await page.locator('input[type="text"]').press("Enter");
     await this.nextFrame();
   }
 
@@ -140,7 +160,7 @@ export class ComfyPage {
   }
 
   async rightClickCanvas() {
-    await this.page.mouse.click(10, 10, { button: 'right' });
+    await this.page.mouse.click(10, 10, { button: "right" });
     await this.nextFrame();
   }
 
@@ -153,7 +173,7 @@ export class ComfyPage {
     await this.canvas.click({
       position: {
         x: 724,
-        y: 625
+        y: 625,
       },
     });
     this.page.mouse.move(10, 10);
@@ -164,34 +184,34 @@ export class ComfyPage {
     await this.canvas.click({
       position: {
         x: 724,
-        y: 645
+        y: 645,
       },
-      button: 'right'
+      button: "right",
     });
     this.page.mouse.move(10, 10);
     await this.nextFrame();
   }
 
   async select2Nodes() {
     // Select 2 CLIP nodes.
-    await this.page.keyboard.down('Control');
+    await this.page.keyboard.down("Control");
     await this.clickTextEncodeNode1();
     await this.clickTextEncodeNode2();
-    await this.page.keyboard.up('Control');
+    await this.page.keyboard.up("Control");
     await this.nextFrame();
   }
 
   async ctrlC() {
-    await this.page.keyboard.down('Control');
-    await this.page.keyboard.press('KeyC');
-    await this.page.keyboard.up('Control');
+    await this.page.keyboard.down("Control");
+    await this.page.keyboard.press("KeyC");
+    await this.page.keyboard.up("Control");
     await this.nextFrame();
   }
 
   async ctrlV() {
-    await this.page.keyboard.down('Control');
-    await this.page.keyboard.press('KeyV');
-    await this.page.keyboard.up('Control');
+    await this.page.keyboard.down("Control");
+    await this.page.keyboard.press("KeyV");
+    await this.page.keyboard.up("Control");
     await this.nextFrame();
   }
 }
@@ -202,25 +222,27 @@ export const comfyPageFixture = base.extend<{ comfyPage: ComfyPage }>({
     await comfyPage.goto();
     // Unify font for consistent screenshots.
     await page.addStyleTag({
-      url: "https://fonts.googleapis.com/css2?family=Roboto:ital,wght@0,100;0,300;0,400;0,500;0,700;0,900;1,100;1,300;1,400;1,500;1,700;1,900&display=swap"
+      url: "https://fonts.googleapis.com/css2?family=Roboto:ital,wght@0,100;0,300;0,400;0,500;0,700;0,900;1,100;1,300;1,400;1,500;1,700;1,900&display=swap",
     });
     await page.addStyleTag({
-      url: "https://fonts.googleapis.com/css2?family=Noto+Color+Emoji&family=Roboto+Mono:ital,wght@0,100..700;1,100..700&family=Roboto:ital,wght@0,100;0,300;0,400;0,500;0,700;0,900;1,100;1,300;1,400;1,500;1,700;1,900&display=swap"
+      url: "https://fonts.googleapis.com/css2?family=Noto+Color+Emoji&family=Roboto+Mono:ital,wght@0,100..700;1,100..700&family=Roboto:ital,wght@0,100;0,300;0,400;0,500;0,700;0,900;1,100;1,300;1,400;1,500;1,700;1,900&display=swap",
     });
     await page.addStyleTag({
       content: `
       * {
 				font-family: 'Roboto Mono', 'Noto Color Emoji';
-			}`
+			}`,
     });
 
     await page.waitForFunction(() => document.fonts.ready);
-    await page.waitForFunction(() => window['app'] != undefined);
-    await page.evaluate(() => { window['app']['canvas'].show_info = false; });
+    await page.waitForFunction(() => window["app"] != undefined);
+    await page.evaluate(() => {
+      window["app"]["canvas"].show_info = false;
+    });
     await comfyPage.nextFrame();
     // Reset view to force re-rendering of canvas. So that info fields like fps
     // become hidden.
     await comfyPage.resetView();
     await use(comfyPage);
   },
-});
+});

browser_tests/interaction.spec.ts-snapshots/disconnected-edge-with-menu-chromium-win32.png

@@ -0,0 +1,47 @@
+import { expect } from "@playwright/test";
+import { ComfyPage, comfyPageFixture } from "./ComfyPage";
+
+export const test = comfyPageFixture.extend<{ comfyPage: ComfyPage }>({
+  comfyPage: async ({ comfyPage }, use) => {
+    await comfyPage.page.evaluate(async () => {
+      await window["app"].ui.settings.setSettingValueAsync(
+        "Comfy.NodeSearchBoxImpl",
+        "default"
+      );
+    });
+    await use(comfyPage);
+  },
+});
+
+test.describe("Node search box", () => {
+  test("Can trigger on empty canvas double click", async ({ comfyPage }) => {
+    await comfyPage.doubleClickCanvas();
+    await expect(comfyPage.searchBox.input).toHaveCount(1);
+  });
+
+  test("Can trigger on link release", async ({ comfyPage }) => {
+    await comfyPage.page.keyboard.down("Shift");
+    await comfyPage.disconnectEdge();
+    await expect(comfyPage.searchBox.input).toHaveCount(1);
+  });
+
+  test("Does not trigger on link release (no shift)", async ({ comfyPage }) => {
+    await comfyPage.disconnectEdge();
+    await expect(comfyPage.searchBox.input).toHaveCount(0);
+  });
+
+  test("Can add node", async ({ comfyPage }) => {
+    await comfyPage.doubleClickCanvas();
+    await expect(comfyPage.searchBox.input).toHaveCount(1);
+    await comfyPage.searchBox.fillAndSelectFirstNode("KSampler");
+    await expect(comfyPage.canvas).toHaveScreenshot("added-node.png");
+  });
+
+  test("Can auto link node", async ({ comfyPage }) => {
+    await comfyPage.page.keyboard.down("Shift");
+    await comfyPage.disconnectEdge();
+    await comfyPage.page.keyboard.up("Shift");
+    await comfyPage.searchBox.fillAndSelectFirstNode("CLIPTextEncode");
+    await expect(comfyPage.canvas).toHaveScreenshot("auto-linked-node.png");
+  });
+});

browser_tests/nodeSearchBox.spec.ts

@@ -36,6 +36,7 @@ export default defineConfig({
     {
       name: 'chromium',
       use: { ...devices['Desktop Chrome'] },
+      timeout: 3000,
     },
 
     // {

browser_tests/nodeSearchBox.spec.ts-snapshots/added-node-chromium-linux.png

@@ -23,7 +23,13 @@ import { app } from "@/scripts/app";
 import { inject, onMounted, onUnmounted, reactive, Ref, ref } from "vue";
 import NodeSearchBox from "./NodeSearchBox.vue";
 import Dialog from "primevue/dialog";
-import { LiteGraph, LiteGraphCanvasEvent } from "@comfyorg/litegraph";
+import {
+  INodeSlot,
+  LiteGraph,
+  LiteGraphCanvasEvent,
+  LGraphNode,
+  LinkReleaseContext,
+} from "@comfyorg/litegraph";
 import {
   FilterAndValue,
   NodeSearchService,
@@ -63,12 +69,47 @@ const closeDialog = () => {
   clearFilters();
   visible.value = false;
 };
+const connectNodeOnLinkRelease = (
+  node: LGraphNode,
+  context: LinkReleaseContext
+) => {
+  const destIsInput = context.node_from !== undefined;
+  const srcNode = (
+    destIsInput ? context.node_from : context.node_to
+  ) as LGraphNode;
+  const srcSlotIndex: number = context.slot_from.slot_index;
+  const linkDataType = destIsInput
+    ? context.type_filter_in
+    : context.type_filter_out;
+  const destSlots = destIsInput ? node.inputs : node.outputs;
+  const destSlotIndex = destSlots.findIndex(
+    (slot: INodeSlot) => slot.type === linkDataType
+  );
+
+  if (destSlotIndex === -1) {
+    console.warn(
+      `Could not find slot with type ${linkDataType} on node ${node.title}`
+    );
+    return;
+  }
+
+  if (destIsInput) {
+    srcNode.connect(srcSlotIndex, node, destSlotIndex);
+  } else {
+    node.connect(destSlotIndex, srcNode, srcSlotIndex);
+  }
+};
 const addNode = (nodeDef: ComfyNodeDef) => {
   closeDialog();
   const node = LiteGraph.createNode(nodeDef.name, nodeDef.display_name, {});
   if (node) {
     node.pos = getNewNodeLocation();
     app.graph.add(node);
+
+    const eventDetail = triggerEvent.value.detail;
+    if (eventDetail.subType === "empty-release") {
+      connectNodeOnLinkRelease(node, eventDetail.linkReleaseContext);
+    }
   }
 };
 const nodeSearchService = (