From 44b109a449e64a4c1b81a1ccc545364b7bfa34de Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Mon, 28 Oct 2024 17:26:52 -0400
Subject: [PATCH] Add type annotation for missingNodeTypes (#1349)

* Add type annotation for missingNodeTypes

* nit

* nit
---
 .../dialog/content/LoadWorkflowWarning.vue    | 12 ++------
 src/scripts/app.ts                            |  6 ++--
 src/services/dialogService.ts                 |  3 +-
 src/types/comfy.d.ts                          | 30 +++++++++++++++++++
 4 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/src/components/dialog/content/LoadWorkflowWarning.vue b/src/components/dialog/content/LoadWorkflowWarning.vue
index 3274c77c3..cc440611e 100644
--- a/src/components/dialog/content/LoadWorkflowWarning.vue
+++ b/src/components/dialog/content/LoadWorkflowWarning.vue
@@ -37,18 +37,10 @@ import { computed } from 'vue'
 import ListBox from 'primevue/listbox'
 import Button from 'primevue/button'
 import NoResultsPlaceholder from '@/components/common/NoResultsPlaceholder.vue'
-
-interface NodeType {
-  type: string
-  hint?: string
-  action?: {
-    text: string
-    callback: () => void
-  }
-}
+import type { MissingNodeType } from '@/types/comfy'
 
 const props = defineProps<{
-  missingNodeTypes: (string | NodeType)[]
+  missingNodeTypes: MissingNodeType[]
 }>()
 
 const uniqueNodes = computed(() => {
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 122d9a797..c3b2cdad6 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -15,7 +15,7 @@ import { addDomClippingSetting } from './domWidget'
 import { createImageHost, calculateImageGrid } from './ui/imagePreview'
 import { DraggableList } from './ui/draggableList'
 import { applyTextReplacements, addStylesheet } from './utils'
-import type { ComfyExtension } from '@/types/comfy'
+import type { ComfyExtension, MissingNodeType } from '@/types/comfy'
 import {
   type ComfyWorkflowJSON,
   type NodeId,
@@ -2188,7 +2188,7 @@ export class ComfyApp {
     localStorage.setItem('litegrapheditor_clipboard', old)
   }
 
-  #showMissingNodesError(missingNodeTypes) {
+  #showMissingNodesError(missingNodeTypes: MissingNodeType[]) {
     if (useSettingStore().get('Comfy.Workflow.ShowMissingNodesWarning')) {
       showLoadWorkflowWarning({ missingNodeTypes })
     }
@@ -2268,7 +2268,7 @@ export class ComfyApp {
       graphData = validatedGraphData ?? graphData
     }
 
-    const missingNodeTypes = []
+    const missingNodeTypes: MissingNodeType[] = []
     const missingModels = []
     await this.#invokeExtensionsAsync(
       'beforeConfigureGraph',
diff --git a/src/services/dialogService.ts b/src/services/dialogService.ts
index b507afea1..bf3ad35df 100644
--- a/src/services/dialogService.ts
+++ b/src/services/dialogService.ts
@@ -11,9 +11,10 @@ import ExecutionErrorDialogContent from '@/components/dialog/content/ExecutionEr
 import TemplateWorkflowsContent from '@/components/templates/TemplateWorkflowsContent.vue'
 import PromptDialogContent from '@/components/dialog/content/PromptDialogContent.vue'
 import { i18n } from '@/i18n'
+import type { MissingNodeType } from '@/types/comfy'
 
 export function showLoadWorkflowWarning(props: {
-  missingNodeTypes: any[]
+  missingNodeTypes: MissingNodeType[]
   [key: string]: any
 }) {
   const dialogStore = useDialogStore()
diff --git a/src/types/comfy.d.ts b/src/types/comfy.d.ts
index 6f44b5f45..7af4129f9 100644
--- a/src/types/comfy.d.ts
+++ b/src/types/comfy.d.ts
@@ -28,6 +28,18 @@ export type MenuCommandGroup = {
   commands: string[]
 }
 
+export type MissingNodeType =
+  | string
+  // Primarily used by group nodes.
+  | {
+      type: string
+      hint?: string
+      action?: {
+        text: string
+        callback: () => void
+      }
+    }
+
 export interface ComfyExtension {
   /**
    * The name of the extension
@@ -121,6 +133,24 @@ export interface ComfyExtension {
    */
   nodeCreated?(node: LGraphNode, app: ComfyApp): void
 
+  /**
+   * Allows the extension to modify the graph data before it is configured.
+   * @param graphData The graph data
+   * @param missingNodeTypes The missing node types
+   */
+  beforeConfigureGraph?(
+    graphData: ComfyWorkflowJSON,
+    missingNodeTypes: MissingNodeType[]
+  ): Promise<void> | void
+
+  /**
+   * Allows the extension to run code after the graph is configured.
+   * @param missingNodeTypes The missing node types
+   */
+  afterConfigureGraph?(
+    missingNodeTypes: MissingNodeType[]
+  ): Promise<void> | void
+
   [key: string]: any
 }