From ebc71b0e46137cab2c3944fb6b42c886ac536db8 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Fri, 4 Oct 2024 15:33:27 -0400
Subject: [PATCH] Add PromptDialog to replace window.prompt (#1104)

* Save file prompt dialog

* Don't download if dialog dismissed

* refactor

* style dialog

* nit

* Autofocus
---
 src/App.vue                                   |  2 +-
 src/components/dialog/GlobalDialog.vue        |  5 +-
 .../content/ExecutionErrorDialogContent.vue   |  4 --
 .../dialog/content/PromptDialogContent.vue    | 43 ++++++++++++++++
 .../templates/TemplateWorkflowsContent.vue    |  2 +-
 src/scripts/app.ts                            |  6 +--
 src/scripts/ui.ts                             | 51 ++-----------------
 src/scripts/ui/menu/index.ts                  | 23 ++++++---
 src/services/dialogService.ts                 | 42 ++++++++++++++-
 src/stores/commandStore.ts                    |  8 +++
 src/stores/dialogStore.ts                     | 28 +++++++---
 tests-ui/globalSetup.ts                       | 10 +++-
 12 files changed, 148 insertions(+), 76 deletions(-)
 create mode 100644 src/components/dialog/content/PromptDialogContent.vue

diff --git a/src/App.vue b/src/App.vue
index 344e8df70..675f4fc87 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -10,7 +10,7 @@
 
 <script setup lang="ts">
 import config from '@/config'
-import { computed, onMounted, onBeforeUnmount } from 'vue'
+import { computed, onMounted } from 'vue'
 import { useWorkspaceStore } from '@/stores/workspaceStateStore'
 import BlockUI from 'primevue/blockui'
 import ProgressSpinner from 'primevue/progressspinner'
diff --git a/src/components/dialog/GlobalDialog.vue b/src/components/dialog/GlobalDialog.vue
index d0359551c..3a1615fbe 100644
--- a/src/components/dialog/GlobalDialog.vue
+++ b/src/components/dialog/GlobalDialog.vue
@@ -11,6 +11,7 @@
     @hide="dialogStore.closeDialog"
     @maximize="onMaximize"
     @unmaximize="onUnmaximize"
+    :pt="{ header: 'pb-0' }"
   >
     <template #header>
       <component
@@ -30,7 +31,9 @@ import { useDialogStore } from '@/stores/dialogStore'
 import Dialog from 'primevue/dialog'
 
 const dialogStore = useDialogStore()
-const maximizable = computed(() => dialogStore.props.maximizable ?? false)
+const maximizable = computed(
+  () => dialogStore.dialogComponentProps.maximizable ?? false
+)
 const maximized = ref(false)
 
 const onMaximize = () => {
diff --git a/src/components/dialog/content/ExecutionErrorDialogContent.vue b/src/components/dialog/content/ExecutionErrorDialogContent.vue
index 5d63bab7d..6dd589609 100644
--- a/src/components/dialog/content/ExecutionErrorDialogContent.vue
+++ b/src/components/dialog/content/ExecutionErrorDialogContent.vue
@@ -191,8 +191,4 @@ const openNewGithubIssue = async () => {
   white-space: pre-wrap;
   word-wrap: break-word;
 }
-
-.no-results-placeholder {
-  padding-top: 0;
-}
 </style>
diff --git a/src/components/dialog/content/PromptDialogContent.vue b/src/components/dialog/content/PromptDialogContent.vue
new file mode 100644
index 000000000..c014ada69
--- /dev/null
+++ b/src/components/dialog/content/PromptDialogContent.vue
@@ -0,0 +1,43 @@
+<template>
+  <div class="prompt-dialog-content flex flex-col gap-2 pt-8">
+    <FloatLabel>
+      <InputText
+        ref="inputRef"
+        v-model="inputValue"
+        @keyup.enter="onConfirm"
+        @focus="selectAllText"
+        autofocus
+      />
+      <label>{{ message }}</label>
+    </FloatLabel>
+    <Button @click="onConfirm">{{ $t('confirm') }}</Button>
+  </div>
+</template>
+
+<script setup lang="ts">
+import Button from 'primevue/button'
+import InputText from 'primevue/inputtext'
+import FloatLabel from 'primevue/floatlabel'
+import { ref } from 'vue'
+import { useDialogStore } from '@/stores/dialogStore'
+
+const props = defineProps<{
+  message: string
+  defaultValue: string
+  onConfirm: (value: string) => void
+}>()
+
+const inputValue = ref<string>(props.defaultValue)
+
+const onConfirm = () => {
+  props.onConfirm(inputValue.value)
+  useDialogStore().closeDialog()
+}
+
+const inputRef = ref(null)
+const selectAllText = () => {
+  if (!inputRef.value) return
+  const inputElement = inputRef.value.$el
+  inputElement.setSelectionRange(0, inputElement.value.length)
+}
+</script>
diff --git a/src/components/templates/TemplateWorkflowsContent.vue b/src/components/templates/TemplateWorkflowsContent.vue
index bbaf67863..65adbb396 100644
--- a/src/components/templates/TemplateWorkflowsContent.vue
+++ b/src/components/templates/TemplateWorkflowsContent.vue
@@ -1,6 +1,6 @@
 <template>
   <div
-    class="flex flex-wrap content-around justify-around gap-4"
+    class="flex flex-wrap content-around justify-around gap-4 mt-4"
     data-testid="template-workflows-content"
   >
     <div
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 7b8a014c9..9eb32ac89 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -2231,8 +2231,7 @@ export class ComfyApp {
     if (useSettingStore().get('Comfy.Workflow.ShowMissingNodesWarning')) {
       showLoadWorkflowWarning({
         missingNodeTypes,
-        hasAddedNodes,
-        maximizable: true
+        hasAddedNodes
       })
     }
 
@@ -2245,8 +2244,7 @@ export class ComfyApp {
     if (useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')) {
       showMissingModelsWarning({
         missingModels,
-        paths,
-        maximizable: true
+        paths
       })
     }
 
diff --git a/src/scripts/ui.ts b/src/scripts/ui.ts
index b121e0fc2..dc565d020 100644
--- a/src/scripts/ui.ts
+++ b/src/scripts/ui.ts
@@ -6,6 +6,7 @@ import { ComfyApp, app } from './app'
 import { TaskItem } from '@/types/apiTypes'
 import { showSettingsDialog } from '@/services/dialogService'
 import { useSettingStore } from '@/stores/settingStore'
+import { useCommandStore } from '@/stores/commandStore'
 
 export const ComfyDialog = _ComfyDialog
 
@@ -585,30 +586,7 @@ export class ComfyUI {
         id: 'comfy-save-button',
         textContent: 'Save',
         onclick: () => {
-          let filename = 'workflow.json'
-          if (useSettingStore().get('Comfy.PromptFilename')) {
-            filename = prompt('Save workflow as:', filename)
-            if (!filename) return
-            if (!filename.toLowerCase().endsWith('.json')) {
-              filename += '.json'
-            }
-          }
-          app.graphToPrompt().then((p) => {
-            const json = JSON.stringify(p.workflow, null, 2) // convert the data to a JSON string
-            const blob = new Blob([json], { type: 'application/json' })
-            const url = URL.createObjectURL(blob)
-            const a = $el('a', {
-              href: url,
-              download: filename,
-              style: { display: 'none' },
-              parent: document.body
-            })
-            a.click()
-            setTimeout(function () {
-              a.remove()
-              window.URL.revokeObjectURL(url)
-            }, 0)
-          })
+          useCommandStore().execute('Comfy.ExportWorkflow')
         }
       }),
       $el('button', {
@@ -616,30 +594,7 @@ export class ComfyUI {
         textContent: 'Save (API Format)',
         style: { width: '100%', display: 'none' },
         onclick: () => {
-          let filename = 'workflow_api.json'
-          if (useSettingStore().get('Comfy.PromptFilename')) {
-            filename = prompt('Save workflow (API) as:', filename)
-            if (!filename) return
-            if (!filename.toLowerCase().endsWith('.json')) {
-              filename += '.json'
-            }
-          }
-          app.graphToPrompt().then((p) => {
-            const json = JSON.stringify(p.output, null, 2) // convert the data to a JSON string
-            const blob = new Blob([json], { type: 'application/json' })
-            const url = URL.createObjectURL(blob)
-            const a = $el('a', {
-              href: url,
-              download: filename,
-              style: { display: 'none' },
-              parent: document.body
-            })
-            a.click()
-            setTimeout(function () {
-              a.remove()
-              window.URL.revokeObjectURL(url)
-            }, 0)
-          })
+          useCommandStore().execute('Comfy.ExportWorkflowAPI')
         }
       }),
       $el('button', {
diff --git a/src/scripts/ui/menu/index.ts b/src/scripts/ui/menu/index.ts
index 0e708e443..17dd88af3 100644
--- a/src/scripts/ui/menu/index.ts
+++ b/src/scripts/ui/menu/index.ts
@@ -2,6 +2,8 @@ import type { ComfyApp } from '@/scripts/app'
 import { $el } from '../../ui'
 import { downloadBlob } from '../../utils'
 import { ComfyButtonGroup } from '../components/buttonGroup'
+import { showPromptDialog } from '@/services/dialogService'
+import { useSettingStore } from '@/stores/settingStore'
 import './menu.css'
 
 // Export to make sure following components are shimmed and exported by vite
@@ -32,13 +34,18 @@ export class ComfyAppMenu {
     ])
   }
 
-  getFilename(defaultName: string) {
-    if (this.app.ui.settings.getSettingValue('Comfy.PromptFilename', true)) {
-      defaultName = prompt('Save workflow as:', defaultName)
-      if (!defaultName) return
-      if (!defaultName.toLowerCase().endsWith('.json')) {
-        defaultName += '.json'
+  async getFilename(defaultName: string): Promise<string | null> {
+    if (useSettingStore().get('Comfy.PromptFilename')) {
+      let filename = await showPromptDialog({
+        title: 'Export Workflow',
+        message: 'Enter the filename:',
+        defaultValue: defaultName
+      })
+      if (!filename) return null
+      if (!filename.toLowerCase().endsWith('.json')) {
+        filename += '.json'
       }
+      return filename
     }
     return defaultName
   }
@@ -46,14 +53,14 @@ export class ComfyAppMenu {
   async exportWorkflow(
     filename: string,
     promptProperty: 'workflow' | 'output'
-  ) {
+  ): Promise<void> {
     if (this.app.workflowManager.activeWorkflow?.path) {
       filename = this.app.workflowManager.activeWorkflow.name
     }
     const p = await this.app.graphToPrompt()
     const json = JSON.stringify(p[promptProperty], null, 2)
     const blob = new Blob([json], { type: 'application/json' })
-    const file = this.getFilename(filename)
+    const file = await this.getFilename(filename)
     if (!file) return
     downloadBlob(file, blob)
   }
diff --git a/src/services/dialogService.ts b/src/services/dialogService.ts
index 86cc8fd4a..6bda88f7a 100644
--- a/src/services/dialogService.ts
+++ b/src/services/dialogService.ts
@@ -9,6 +9,7 @@ import SettingDialogHeader from '@/components/dialog/header/SettingDialogHeader.
 import type { ExecutionErrorWsMessage } from '@/types/apiTypes'
 import ExecutionErrorDialogContent from '@/components/dialog/content/ExecutionErrorDialogContent.vue'
 import TemplateWorkflowsContent from '@/components/templates/TemplateWorkflowsContent.vue'
+import PromptDialogContent from '@/components/dialog/content/PromptDialogContent.vue'
 import { i18n } from '@/i18n'
 
 export function showLoadWorkflowWarning(props: {
@@ -19,7 +20,10 @@ export function showLoadWorkflowWarning(props: {
   const dialogStore = useDialogStore()
   dialogStore.showDialog({
     component: LoadWorkflowWarning,
-    props
+    props,
+    dialogComponentProps: {
+      maximizable: true
+    }
   })
 }
 
@@ -31,7 +35,10 @@ export function showMissingModelsWarning(props: {
   const dialogStore = useDialogStore()
   dialogStore.showDialog({
     component: MissingModelsWarning,
-    props
+    props,
+    dialogComponentProps: {
+      maximizable: true
+    }
   })
 }
 
@@ -57,3 +64,34 @@ export function showTemplateWorkflowsDialog() {
     component: TemplateWorkflowsContent
   })
 }
+
+export async function showPromptDialog({
+  title,
+  message,
+  defaultValue = ''
+}: {
+  title: string
+  message: string
+  defaultValue?: string
+}): Promise<string | null> {
+  const dialogStore = useDialogStore()
+
+  return new Promise((resolve) => {
+    dialogStore.showDialog({
+      title,
+      component: PromptDialogContent,
+      props: {
+        message,
+        defaultValue,
+        onConfirm: (value: string) => {
+          resolve(value)
+        }
+      },
+      dialogComponentProps: {
+        onClose: () => {
+          resolve(null)
+        }
+      }
+    })
+  })
+}
diff --git a/src/stores/commandStore.ts b/src/stores/commandStore.ts
index 7263d2b9a..918a7bebd 100644
--- a/src/stores/commandStore.ts
+++ b/src/stores/commandStore.ts
@@ -320,6 +320,13 @@ export const useCommandStore = defineStore('command', () => {
     return commandsById.value[command]
   }
 
+  const execute = (commandId: string) => {
+    const command = getCommand(commandId)
+    if (command) {
+      command.function()
+    }
+  }
+
   const isRegistered = (command: string) => {
     return !!commandsById.value[command]
   }
@@ -334,6 +341,7 @@ export const useCommandStore = defineStore('command', () => {
 
   return {
     commands,
+    execute,
     getCommand,
     getCommandFunction,
     registerCommand,
diff --git a/src/stores/dialogStore.ts b/src/stores/dialogStore.ts
index 67ec6dab0..86ba93e32 100644
--- a/src/stores/dialogStore.ts
+++ b/src/stores/dialogStore.ts
@@ -2,14 +2,22 @@
 // Currently we need to bridge between legacy app code and Vue app with a Pinia store.
 
 import { defineStore } from 'pinia'
-import { type Component, markRaw } from 'vue'
+import { type Component, markRaw, nextTick } from 'vue'
 
 interface DialogState {
   isVisible: boolean
   title: string
   headerComponent: Component | null
   component: Component | null
+  // Props passing to the component
   props: Record<string, any>
+  // Props passing to the Dialog component
+  dialogComponentProps: DialogComponentProps
+}
+
+interface DialogComponentProps {
+  maximizable?: boolean
+  onClose?: () => void
 }
 
 export const useDialogStore = defineStore('dialog', {
@@ -18,7 +26,8 @@ export const useDialogStore = defineStore('dialog', {
     title: '',
     headerComponent: null,
     component: null,
-    props: {}
+    props: {},
+    dialogComponentProps: {}
   }),
 
   actions: {
@@ -27,15 +36,22 @@ export const useDialogStore = defineStore('dialog', {
       headerComponent?: Component
       component: Component
       props?: Record<string, any>
+      dialogComponentProps?: DialogComponentProps
     }) {
-      this.title = options.title
-      this.headerComponent = markRaw(options.headerComponent)
-      this.component = markRaw(options.component)
-      this.props = options.props || {}
       this.isVisible = true
+      nextTick(() => {
+        this.title = options.title
+        this.headerComponent = markRaw(options.headerComponent)
+        this.component = markRaw(options.component)
+        this.props = options.props || {}
+        this.dialogComponentProps = options.dialogComponentProps || {}
+      })
     },
 
     closeDialog() {
+      if (this.dialogComponentProps.onClose) {
+        this.dialogComponentProps.onClose()
+      }
       this.isVisible = false
     }
   }
diff --git a/tests-ui/globalSetup.ts b/tests-ui/globalSetup.ts
index 16c6d9c22..b9c81d287 100644
--- a/tests-ui/globalSetup.ts
+++ b/tests-ui/globalSetup.ts
@@ -16,7 +16,15 @@ module.exports = async function () {
   jest.mock('@/services/dialogService', () => {
     return {
       showLoadWorkflowWarning: jest.fn(),
-      showMissingModelsWarning: jest.fn()
+      showMissingModelsWarning: jest.fn(),
+      showSettingsDialog: jest.fn(),
+      showExecutionErrorDialog: jest.fn(),
+      showTemplateWorkflowsDialog: jest.fn(),
+      showPromptDialog: jest
+        .fn()
+        .mockImplementation((message, defaultValue) => {
+          return Promise.resolve(defaultValue)
+        })
     }
   })
 }