From 9cea54686a6d0b7ca3bc1803196865da716601c0 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Fri, 16 Aug 2024 13:22:17 -0400
Subject: [PATCH] Pin reroute on link release searchbox (#469)

* Correctly handle wildcard type

* Add test

* nit

* Pin reroute on link release search

* Connect wildcard

* nit
---
 browser_tests/ComfyPage.ts                     | 10 ++++++++--
 browser_tests/nodeSearchBox.spec.ts            | 10 ++++++++--
 src/components/searchbox/NodeSearchBox.vue     | 18 ++++++++++++------
 .../searchbox/NodeSearchBoxPopover.vue         |  4 ++++
 src/services/nodeSearchService.ts              |  8 +++++++-
 src/types/litegraphTypes.ts                    |  3 ++-
 tests-ui/tests/nodeSearchService.test.ts       |  2 ++
 7 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/browser_tests/ComfyPage.ts b/browser_tests/ComfyPage.ts
index f8f16071..caa8ae92 100644
--- a/browser_tests/ComfyPage.ts
+++ b/browser_tests/ComfyPage.ts
@@ -26,14 +26,20 @@ class ComfyNodeSearchBox {
     )
   }
 
-  async fillAndSelectFirstNode(nodeName: string) {
+  async fillAndSelectFirstNode(
+    nodeName: string,
+    options?: { suggestionIndex: number }
+  ) {
     await this.input.waitFor({ state: 'visible' })
     await this.input.fill(nodeName)
     await this.dropdown.waitFor({ state: 'visible' })
     // Wait for some time for the auto complete list to update.
     // The auto complete list is debounced and may take some time to update.
     await this.page.waitForTimeout(500)
-    await this.dropdown.locator('li').nth(0).click()
+    await this.dropdown
+      .locator('li')
+      .nth(options?.suggestionIndex || 0)
+      .click()
   }
 }
 
diff --git a/browser_tests/nodeSearchBox.spec.ts b/browser_tests/nodeSearchBox.spec.ts
index 0805d532..673a0db9 100644
--- a/browser_tests/nodeSearchBox.spec.ts
+++ b/browser_tests/nodeSearchBox.spec.ts
@@ -35,7 +35,10 @@ test.describe('Node search box', () => {
 
   test('Can auto link node', async ({ comfyPage }) => {
     await comfyPage.disconnectEdge()
-    await comfyPage.searchBox.fillAndSelectFirstNode('CLIPTextEncode')
+    // Select the second item as the first item is always reroute
+    await comfyPage.searchBox.fillAndSelectFirstNode('CLIPTextEncode', {
+      suggestionIndex: 1
+    })
     await expect(comfyPage.canvas).toHaveScreenshot('auto-linked-node.png')
   })
 
@@ -54,7 +57,10 @@ test.describe('Node search box', () => {
     await comfyPage.dragAndDrop(outputSlot1Pos, emptySpacePos)
     await comfyPage.page.keyboard.up('Shift')
 
-    await comfyPage.searchBox.fillAndSelectFirstNode('Load Checkpoint')
+    // Select the second item as the first item is always reroute
+    await comfyPage.searchBox.fillAndSelectFirstNode('Load Checkpoint', {
+      suggestionIndex: 1
+    })
     await expect(comfyPage.canvas).toHaveScreenshot(
       'auto-linked-node-batch.png'
     )
diff --git a/src/components/searchbox/NodeSearchBox.vue b/src/components/searchbox/NodeSearchBox.vue
index 2877b94a..b918a0df 100644
--- a/src/components/searchbox/NodeSearchBox.vue
+++ b/src/components/searchbox/NodeSearchBox.vue
@@ -78,6 +78,11 @@ const props = defineProps({
   searchLimit: {
     type: Number,
     default: 64
+  },
+  // TODO: Find a more flexible mechanism to add pinned nodes
+  includeReroute: {
+    type: Boolean,
+    default: false
   }
 })
 
@@ -91,13 +96,14 @@ const placeholder = computed(() => {
 
 const search = (query: string) => {
   currentQuery.value = query
-  suggestions.value = useNodeDefStore().nodeSearchService.searchNode(
-    query,
-    props.filters,
-    {
+  suggestions.value = [
+    ...(props.includeReroute
+      ? [useNodeDefStore().nodeDefsByName['Reroute']]
+      : []),
+    ...useNodeDefStore().nodeSearchService.searchNode(query, props.filters, {
       limit: props.searchLimit
-    }
-  )
+    })
+  ]
 }
 
 const highlightQuery = (text: string, query: string) => {
diff --git a/src/components/searchbox/NodeSearchBoxPopover.vue b/src/components/searchbox/NodeSearchBoxPopover.vue
index 7ae1928a..5cc1ae7a 100644
--- a/src/components/searchbox/NodeSearchBoxPopover.vue
+++ b/src/components/searchbox/NodeSearchBoxPopover.vue
@@ -20,6 +20,7 @@
       <template #container>
         <NodeSearchBox
           :filters="nodeFilters"
+          :includeReroute="includeReroute"
           @add-filter="addFilter"
           @remove-filter="removeFilter"
           @add-node="addNode"
@@ -61,6 +62,7 @@ const getNewNodeLocation = (): [number, number] => {
   return [originalEvent.canvasX, originalEvent.canvasY]
 }
 const nodeFilters = reactive([])
+const includeReroute = ref(false)
 const addFilter = (filter: FilterAndValue) => {
   nodeFilters.push(filter)
 }
@@ -100,6 +102,7 @@ const linkReleaseTriggerMode = computed(() => {
 })
 
 const canvasEventHandler = (e: LiteGraphCanvasEvent) => {
+  includeReroute.value = false
   const shiftPressed = (e.detail.originalEvent as KeyboardEvent).shiftKey
 
   if (e.detail.subType === 'empty-release') {
@@ -123,6 +126,7 @@ const canvasEventHandler = (e: LiteGraphCanvasEvent) => {
     )
     const dataType = firstLink.type
     addFilter([filter, dataType])
+    includeReroute.value = true
   }
   triggerEvent.value = e
   visible.value = true
diff --git a/src/services/nodeSearchService.ts b/src/services/nodeSearchService.ts
index 0840e46d..976689c0 100644
--- a/src/services/nodeSearchService.ts
+++ b/src/services/nodeSearchService.ts
@@ -57,7 +57,13 @@ export abstract class NodeFilter<FilterOptionT = string> {
   public abstract getNodeOptions(node: ComfyNodeDefImpl): FilterOptionT[]
 
   public matches(node: ComfyNodeDefImpl, value: FilterOptionT): boolean {
-    return this.getNodeOptions(node).includes(value)
+    if (value === '*') {
+      return true
+    }
+    const options = this.getNodeOptions(node)
+    return (
+      options.includes(value) || _.some(options, (option) => option === '*')
+    )
   }
 }
 
diff --git a/src/types/litegraphTypes.ts b/src/types/litegraphTypes.ts
index d2089dc2..6325badf 100644
--- a/src/types/litegraphTypes.ts
+++ b/src/types/litegraphTypes.ts
@@ -55,7 +55,8 @@ export class ConnectingLinkImpl implements ConnectingLink {
     const newNodeSlots =
       this.releaseSlotType === 'output' ? newNode.outputs : newNode.inputs
     const newNodeSlot = newNodeSlots.findIndex(
-      (slot: INodeSlot) => slot.type === this.type
+      (slot: INodeSlot) =>
+        slot.type === this.type || slot.type === '*' || this.type === '*'
     )
 
     if (newNodeSlot === -1) {
diff --git a/tests-ui/tests/nodeSearchService.test.ts b/tests-ui/tests/nodeSearchService.test.ts
index 85384361..14c68965 100644
--- a/tests-ui/tests/nodeSearchService.test.ts
+++ b/tests-ui/tests/nodeSearchService.test.ts
@@ -58,6 +58,8 @@ describe('nodeSearchService', () => {
     const service = new NodeSearchService(EXAMPLE_NODE_DEFS)
     const inputFilter = service.getFilterById('input')
     expect(service.searchNode('L', [[inputFilter, 'LATENT']])).toHaveLength(1)
+    // Wildcard should match all.
+    expect(service.searchNode('L', [[inputFilter, '*']])).toHaveLength(2)
     expect(service.searchNode('L')).toHaveLength(2)
   })
 })