From 9d855d637e65f409ca4cb867d2fe8f530ffab635 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Thu, 5 Sep 2024 08:00:41 -0700
Subject: [PATCH] Ignore missing dialog (#743)

* Do not report missing nodes/models warning repeatedly

* Add playwright tests

* cast finalOptions, add comments to interface

* Use old menu in tests to not break top left click methods

* Assert no dialog on undo and on redo separately

* nit

* nit

---------

Co-authored-by: christian-byrne <abolkonsky.rem@gmail.com>
---
 browser_tests/ComfyPage.ts   |  7 ++++++
 browser_tests/dialog.spec.ts | 42 ++++++++++++++++++++++++++++++++++++
 src/scripts/app.ts           |  7 +++---
 src/scripts/changeTracker.ts |  5 ++++-
 src/scripts/workflows.ts     |  6 +++++-
 5 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/browser_tests/ComfyPage.ts b/browser_tests/ComfyPage.ts
index 7a8f2d90d..96d89dff4 100644
--- a/browser_tests/ComfyPage.ts
+++ b/browser_tests/ComfyPage.ts
@@ -424,6 +424,13 @@ export class ComfyPage {
     await this.nextFrame()
   }
 
+  async ctrlY() {
+    await this.page.keyboard.down('Control')
+    await this.page.keyboard.press('KeyY')
+    await this.page.keyboard.up('Control')
+    await this.nextFrame()
+  }
+
   async ctrlArrowUp() {
     await this.page.keyboard.down('Control')
     await this.page.keyboard.press('ArrowUp')
diff --git a/browser_tests/dialog.spec.ts b/browser_tests/dialog.spec.ts
index f23ef284b..a869407ba 100644
--- a/browser_tests/dialog.spec.ts
+++ b/browser_tests/dialog.spec.ts
@@ -13,6 +13,48 @@ test.describe('Load workflow warning', () => {
   })
 })
 
+test('Does not report warning when switching between opened workflows', async ({
+  comfyPage
+}) => {
+  await comfyPage.setSetting('Comfy.UseNewMenu', 'Top')
+  await comfyPage.loadWorkflow('missing_nodes')
+  await comfyPage.page.locator('.p-dialog-close-button').click()
+
+  // Load default workflow
+  const workflowSelector = comfyPage.page.locator(
+    'button.comfyui-workflows-button'
+  )
+  await workflowSelector.hover()
+  await workflowSelector.click()
+  await comfyPage.page.locator('button[title="Load default workflow"]').click()
+
+  // Switch back to the missing_nodes workflow
+  await workflowSelector.click()
+  await comfyPage.page.locator('span:has-text("missing_nodes")').first().click()
+  await comfyPage.nextFrame()
+
+  await expect(comfyPage.page.locator('.comfy-missing-nodes')).not.toBeVisible()
+  await comfyPage.setSetting('Comfy.UseNewMenu', 'Disabled')
+})
+
+test('Does not report warning on undo/redo', async ({ comfyPage }) => {
+  await comfyPage.loadWorkflow('missing_nodes')
+  await comfyPage.page.locator('.p-dialog-close-button').click()
+  await comfyPage.nextFrame()
+
+  // Make a change to the graph
+  await comfyPage.setSetting('Comfy.NodeSearchBoxImpl', 'default')
+  await comfyPage.page.waitForTimeout(256)
+  await comfyPage.doubleClickCanvas()
+  await comfyPage.searchBox.fillAndSelectFirstNode('KSampler')
+
+  // Undo and redo the change
+  await comfyPage.ctrlZ()
+  await expect(comfyPage.page.locator('.comfy-missing-nodes')).not.toBeVisible()
+  await comfyPage.ctrlY()
+  await expect(comfyPage.page.locator('.comfy-missing-nodes')).not.toBeVisible()
+})
+
 test.describe('Execution error', () => {
   test('Should display an error message when an execution error occurs', async ({
     comfyPage
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 8729f816b..faee493cf 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -2206,7 +2206,8 @@ export class ComfyApp {
     graphData?: ComfyWorkflowJSON,
     clean: boolean = true,
     restore_view: boolean = true,
-    workflow: string | null | ComfyWorkflow = null
+    workflow: string | null | ComfyWorkflow = null,
+    { showMissingNodesDialog = true, showMissingModelsDialog = true } = {}
   ) {
     if (clean !== false) {
       this.clean()
@@ -2398,10 +2399,10 @@ export class ComfyApp {
     }
 
     // TODO: Properly handle if both nodes and models are missing (sequential dialogs?)
-    if (missingNodeTypes.length) {
+    if (missingNodeTypes.length && showMissingNodesDialog) {
       this.showMissingNodesError(missingNodeTypes)
     }
-    if (missingModels.length) {
+    if (missingModels.length && showMissingModelsDialog) {
       this.showMissingModelsError(missingModels)
     }
     await this.#invokeExtensionsAsync('afterConfigureGraph', missingNodeTypes)
diff --git a/src/scripts/changeTracker.ts b/src/scripts/changeTracker.ts
index e4a4ddc07..210cfd2cf 100644
--- a/src/scripts/changeTracker.ts
+++ b/src/scripts/changeTracker.ts
@@ -72,7 +72,10 @@ export class ChangeTracker {
     if (prevState) {
       target.push(this.activeState)
       this.isOurLoad = true
-      await this.app.loadGraphData(prevState, false, false, this.workflow)
+      await this.app.loadGraphData(prevState, false, false, this.workflow, {
+        showMissingModelsDialog: false,
+        showMissingNodesDialog: false
+      })
       this.activeState = prevState
     }
   }
diff --git a/src/scripts/workflows.ts b/src/scripts/workflows.ts
index ecd77f6a5..b764f40f9 100644
--- a/src/scripts/workflows.ts
+++ b/src/scripts/workflows.ts
@@ -305,7 +305,11 @@ export class ComfyWorkflow {
         this.changeTracker.activeState,
         true,
         true,
-        this
+        this,
+        {
+          showMissingModelsDialog: false,
+          showMissingNodesDialog: false
+        }
       )
     } else {
       const data = await this.getWorkflowData()