From fa75614dc3bdebffb91bf74be9fd45ee94f3c213 Mon Sep 17 00:00:00 2001
From: Benjamin Lu <benceruleanlu@proton.me>
Date: Sun, 6 Apr 2025 18:48:00 -0400
Subject: [PATCH] Add autosave feature (#3330)

Co-authored-by: Benjamin Lu <templu1107@proton.me>
---
 src/components/graph/GraphCanvas.vue          |   2 +
 src/components/topbar/WorkflowTab.vue         |  37 ++-
 src/composables/useWorkflowAutoSave.ts        |  95 ++++++
 src/constants/coreSettings.ts                 |  16 +
 src/schemas/apiSchema.ts                      |   2 +
 .../composables/useWorkflowAutoSave.test.ts   | 274 ++++++++++++++++++
 6 files changed, 416 insertions(+), 10 deletions(-)
 create mode 100644 src/composables/useWorkflowAutoSave.ts
 create mode 100644 tests-ui/tests/composables/useWorkflowAutoSave.test.ts

diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue
index a3e7e9a64..e25dd44cd 100644
--- a/src/components/graph/GraphCanvas.vue
+++ b/src/components/graph/GraphCanvas.vue
@@ -59,6 +59,7 @@ import { useCopy } from '@/composables/useCopy'
 import { useGlobalLitegraph } from '@/composables/useGlobalLitegraph'
 import { useLitegraphSettings } from '@/composables/useLitegraphSettings'
 import { usePaste } from '@/composables/usePaste'
+import { useWorkflowAutoSave } from '@/composables/useWorkflowAutoSave'
 import { useWorkflowPersistence } from '@/composables/useWorkflowPersistence'
 import { CORE_SETTINGS } from '@/constants/coreSettings'
 import { i18n } from '@/i18n'
@@ -233,6 +234,7 @@ onMounted(async () => {
   useContextMenuTranslation()
   useCopy()
   usePaste()
+  useWorkflowAutoSave()
 
   comfyApp.vueAppReady = true
 
diff --git a/src/components/topbar/WorkflowTab.vue b/src/components/topbar/WorkflowTab.vue
index e996be69b..dc0bfa929 100644
--- a/src/components/topbar/WorkflowTab.vue
+++ b/src/components/topbar/WorkflowTab.vue
@@ -7,15 +7,7 @@
       {{ workflowOption.workflow.filename }}
     </span>
     <div class="relative">
-      <span
-        class="status-indicator"
-        v-if="
-          !workspaceStore.shiftDown &&
-          (workflowOption.workflow.isModified ||
-            !workflowOption.workflow.isPersisted)
-        "
-        >•</span
-      >
+      <span class="status-indicator" v-if="shouldShowStatusIndicator">•</span>
       <Button
         class="close-button p-0 w-auto"
         icon="pi pi-times"
@@ -30,7 +22,7 @@
 
 <script setup lang="ts">
 import Button from 'primevue/button'
-import { ref } from 'vue'
+import { computed, ref } from 'vue'
 import { useI18n } from 'vue-i18n'
 
 import {
@@ -38,6 +30,7 @@ import {
   usePragmaticDroppable
 } from '@/composables/usePragmaticDragAndDrop'
 import { useWorkflowService } from '@/services/workflowService'
+import { useSettingStore } from '@/stores/settingStore'
 import { ComfyWorkflow } from '@/stores/workflowStore'
 import { useWorkflowStore } from '@/stores/workflowStore'
 import { useWorkspaceStore } from '@/stores/workspaceStore'
@@ -56,8 +49,32 @@ const { t } = useI18n()
 
 const workspaceStore = useWorkspaceStore()
 const workflowStore = useWorkflowStore()
+const settingStore = useSettingStore()
 const workflowTabRef = ref<HTMLElement | null>(null)
 
+// Use computed refs to cache autosave settings
+const autoSaveSetting = computed(() =>
+  settingStore.get('Comfy.Workflow.AutoSave')
+)
+const autoSaveDelay = computed(() =>
+  settingStore.get('Comfy.Workflow.AutoSaveDelay')
+)
+
+const shouldShowStatusIndicator = computed(() => {
+  // Return true if:
+  // 1. The shift key is not pressed (hence no override).
+  // 2. The workflow is either modified or not yet persisted.
+  // 3. AutoSave is either turned off, or set to 'after delay'
+  //    with a delay longer than 3000ms.
+  return (
+    !workspaceStore.shiftDown &&
+    (props.workflowOption.workflow.isModified ||
+      !props.workflowOption.workflow.isPersisted) &&
+    (autoSaveSetting.value === 'off' ||
+      (autoSaveSetting.value === 'after delay' && autoSaveDelay.value > 3000))
+  )
+})
+
 const closeWorkflows = async (options: WorkflowOption[]) => {
   for (const opt of options) {
     if (
diff --git a/src/composables/useWorkflowAutoSave.ts b/src/composables/useWorkflowAutoSave.ts
new file mode 100644
index 000000000..63eb2eb56
--- /dev/null
+++ b/src/composables/useWorkflowAutoSave.ts
@@ -0,0 +1,95 @@
+import { computed, onUnmounted, watch } from 'vue'
+
+import { api } from '@/scripts/api'
+import { useWorkflowService } from '@/services/workflowService'
+import { useSettingStore } from '@/stores/settingStore'
+import { useWorkflowStore } from '@/stores/workflowStore'
+
+export function useWorkflowAutoSave() {
+  const workflowStore = useWorkflowStore()
+  const settingStore = useSettingStore()
+  const workflowService = useWorkflowService()
+
+  // Use computed refs to cache autosave settings
+  const autoSaveSetting = computed(() =>
+    settingStore.get('Comfy.Workflow.AutoSave')
+  )
+  const autoSaveDelay = computed(() =>
+    settingStore.get('Comfy.Workflow.AutoSaveDelay')
+  )
+
+  let autoSaveTimeout: NodeJS.Timeout | null = null
+  let isSaving = false
+  let needsAutoSave = false
+
+  const scheduleAutoSave = () => {
+    // Clear any existing timeout
+    if (autoSaveTimeout) {
+      clearTimeout(autoSaveTimeout)
+      autoSaveTimeout = null
+    }
+
+    // If autosave is enabled
+    if (autoSaveSetting.value === 'after delay') {
+      // If a save is in progress, mark that we need an autosave after saving
+      if (isSaving) {
+        needsAutoSave = true
+        return
+      }
+      const delay = autoSaveDelay.value
+      autoSaveTimeout = setTimeout(async () => {
+        const activeWorkflow = workflowStore.activeWorkflow
+        if (activeWorkflow?.isModified) {
+          try {
+            isSaving = true
+            await workflowService.saveWorkflow(activeWorkflow)
+          } catch (err) {
+            console.error('Auto save failed:', err)
+          } finally {
+            isSaving = false
+            if (needsAutoSave) {
+              needsAutoSave = false
+              scheduleAutoSave()
+            }
+          }
+        }
+      }, delay)
+    }
+  }
+
+  // Watch for autosave setting changes
+  watch(
+    autoSaveSetting,
+    (newSetting) => {
+      // Clear any existing timeout when settings change
+      if (autoSaveTimeout) {
+        clearTimeout(autoSaveTimeout)
+        autoSaveTimeout = null
+      }
+
+      // If there's an active modified workflow and autosave is enabled, schedule a save
+      if (
+        newSetting === 'after delay' &&
+        workflowStore.activeWorkflow?.isModified
+      ) {
+        scheduleAutoSave()
+      }
+    },
+    { immediate: true }
+  )
+
+  // Listen for graph changes and schedule autosave when they occur
+  const onGraphChanged = () => {
+    scheduleAutoSave()
+  }
+
+  api.addEventListener('graphChanged', onGraphChanged)
+
+  onUnmounted(() => {
+    if (autoSaveTimeout) {
+      clearTimeout(autoSaveTimeout)
+      autoSaveTimeout = null
+    }
+    api.removeEventListener('graphChanged', onGraphChanged)
+  })
+}
diff --git a/src/constants/coreSettings.ts b/src/constants/coreSettings.ts
index dc8e96102..28a02de1a 100644
--- a/src/constants/coreSettings.ts
+++ b/src/constants/coreSettings.ts
@@ -771,5 +771,21 @@ export const CORE_SETTINGS: SettingParams[] = [
     type: 'hidden',
     defaultValue: false,
     versionAdded: '1.15.12'
+  },
+  {
+    id: 'Comfy.Workflow.AutoSaveDelay',
+    name: 'Auto Save Delay (ms)',
+    defaultValue: 1000,
+    type: 'number',
+    tooltip: 'Only applies if Auto Save is set to "after delay".',
+    versionAdded: '1.16.0'
+  },
+  {
+    id: 'Comfy.Workflow.AutoSave',
+    name: 'Auto Save',
+    type: 'combo',
+    options: ['off', 'after delay'], // Room for other options like on focus change, tab change, window change
+    defaultValue: 'off', // Popular requst by users (https://github.com/Comfy-Org/ComfyUI_frontend/issues/1584#issuecomment-2536610154)
+    versionAdded: '1.16.0'
   }
 ]
diff --git a/src/schemas/apiSchema.ts b/src/schemas/apiSchema.ts
index 873d54e3c..cb741129d 100644
--- a/src/schemas/apiSchema.ts
+++ b/src/schemas/apiSchema.ts
@@ -420,6 +420,8 @@ const zSettings = z.record(z.any()).and(
       'Comfy.Server.LaunchArgs': z.record(z.string(), z.string()),
       'LiteGraph.Canvas.MaximumFps': z.number(),
       'Comfy.Workflow.ConfirmDelete': z.boolean(),
+      'Comfy.Workflow.AutoSaveDelay': z.number(),
+      'Comfy.Workflow.AutoSave': z.enum(['off', 'after delay']),
       'Comfy.RerouteBeta': z.boolean(),
       'LiteGraph.Canvas.LowQualityRenderingZoomThreshold': z.number(),
       'Comfy.Canvas.SelectionToolbox': z.boolean()
diff --git a/tests-ui/tests/composables/useWorkflowAutoSave.test.ts b/tests-ui/tests/composables/useWorkflowAutoSave.test.ts
new file mode 100644
index 000000000..2d9ce3ff2
--- /dev/null
+++ b/tests-ui/tests/composables/useWorkflowAutoSave.test.ts
@@ -0,0 +1,274 @@
+import { mount } from '@vue/test-utils'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { useWorkflowAutoSave } from '@/composables/useWorkflowAutoSave'
+import { api } from '@/scripts/api'
+import { useWorkflowService } from '@/services/workflowService'
+
+vi.mock('@/scripts/api', () => ({
+  api: {
+    addEventListener: vi.fn(),
+    removeEventListener: vi.fn()
+  }
+}))
+
+vi.mock('@/services/workflowService', () => ({
+  useWorkflowService: vi.fn(() => ({
+    saveWorkflow: vi.fn()
+  }))
+}))
+
+vi.mock('@/stores/settingStore', () => ({
+  useSettingStore: vi.fn(() => ({
+    get: vi.fn((key) => {
+      if (key === 'Comfy.Workflow.AutoSave') return mockAutoSaveSetting
+      if (key === 'Comfy.Workflow.AutoSaveDelay') return mockAutoSaveDelay
+      return null
+    })
+  }))
+}))
+
+vi.mock('@/stores/workflowStore', () => ({
+  useWorkflowStore: vi.fn(() => ({
+    activeWorkflow: mockActiveWorkflow
+  }))
+}))
+
+let mockAutoSaveSetting: string = 'off'
+let mockAutoSaveDelay: number = 1000
+let mockActiveWorkflow: { isModified: boolean } | null = null
+
+describe('useWorkflowAutoSave', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    vi.useFakeTimers()
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  it('should auto-save workflow after delay when modified and autosave enabled', async () => {
+    mockAutoSaveSetting = 'after delay'
+    mockAutoSaveDelay = 1000
+    mockActiveWorkflow = { isModified: true }
+
+    mount({
+      template: `<div></div>`,
+      setup() {
+        useWorkflowAutoSave()
+        return {}
+      }
+    })
+
+    vi.advanceTimersByTime(1000)
+
+    const serviceInstance = (useWorkflowService as any).mock.results[0].value
+    expect(serviceInstance.saveWorkflow).toHaveBeenCalledWith(
+      mockActiveWorkflow
+    )
+  })
+
+  it('should not auto-save workflow after delay when not modified and autosave enabled', async () => {
+    mockAutoSaveSetting = 'after delay'
+    mockAutoSaveDelay = 1000
+    mockActiveWorkflow = { isModified: false }
+
+    mount({
+      template: `<div></div>`,
+      setup() {
+        useWorkflowAutoSave()
+        return {}
+      }
+    })
+
+    vi.advanceTimersByTime(1000)
+
+    const serviceInstance = (useWorkflowService as any).mock.results[0].value
+    expect(serviceInstance.saveWorkflow).not.toHaveBeenCalledWith(
+      mockActiveWorkflow
+    )
+  })
+
+  it('should not auto save workflow when autosave is off', async () => {
+    mockAutoSaveSetting = 'off'
+    mockAutoSaveDelay = 1000
+    mockActiveWorkflow = { isModified: true }
+
+    mount({
+      template: `<div></div>`,
+      setup() {
+        useWorkflowAutoSave()
+        return {}
+      }
+    })
+
+    vi.advanceTimersByTime(mockAutoSaveDelay)
+
+    const serviceInstance = (useWorkflowService as any).mock.results[0].value
+    expect(serviceInstance.saveWorkflow).not.toHaveBeenCalled()
+  })
+
+  it('should respect the user specified auto save delay', async () => {
+    mockAutoSaveSetting = 'after delay'
+    mockAutoSaveDelay = 2000
+    mockActiveWorkflow = { isModified: true }
+
+    mount({
+      template: `<div></div>`,
+      setup() {
+        useWorkflowAutoSave()
+        return {}
+      }
+    })
+
+    vi.advanceTimersByTime(1000)
+
+    const serviceInstance = (useWorkflowService as any).mock.results[0].value
+    expect(serviceInstance.saveWorkflow).not.toHaveBeenCalled()
+
+    vi.advanceTimersByTime(1000)
+
+    expect(serviceInstance.saveWorkflow).toHaveBeenCalled()
+  })
+
+  it('should debounce save requests', async () => {
+    mockAutoSaveSetting = 'after delay'
+    mockAutoSaveDelay = 2000
+    mockActiveWorkflow = { isModified: true }
+
+    mount({
+      template: `<div></div>`,
+      setup() {
+        useWorkflowAutoSave()
+        return {}
+      }
+    })
+
+    const serviceInstance = (useWorkflowService as any).mock.results[0].value
+    const graphChangedCallback = (api.addEventListener as any).mock.calls[0][1]
+
+    graphChangedCallback()
+
+    vi.advanceTimersByTime(500)
+
+    graphChangedCallback()
+
+    vi.advanceTimersByTime(1999)
+    expect(serviceInstance.saveWorkflow).not.toHaveBeenCalled()
+
+    vi.advanceTimersByTime(1)
+    expect(serviceInstance.saveWorkflow).toHaveBeenCalledTimes(1)
+  })
+
+  it('should handle save error gracefully', async () => {
+    mockAutoSaveSetting = 'after delay'
+    mockAutoSaveDelay = 1000
+    mockActiveWorkflow = { isModified: true }
+
+    const consoleErrorSpy = vi
+      .spyOn(console, 'error')
+      .mockImplementation(() => {})
+
+    try {
+      mount({
+        template: `<div></div>`,
+        setup() {
+          useWorkflowAutoSave()
+          return {}
+        }
+      })
+
+      const serviceInstance = (useWorkflowService as any).mock.results[0].value
+      serviceInstance.saveWorkflow.mockRejectedValue(new Error('Test Error'))
+
+      vi.advanceTimersByTime(1000)
+      await Promise.resolve()
+
+      expect(consoleErrorSpy).toHaveBeenCalledWith(
+        'Auto save failed:',
+        expect.any(Error)
+      )
+    } finally {
+      consoleErrorSpy.mockRestore()
+    }
+  })
+
+  it('should queue autosave requests during saving and reschedule after save completes', async () => {
+    mockAutoSaveSetting = 'after delay'
+    mockAutoSaveDelay = 1000
+    mockActiveWorkflow = { isModified: true }
+
+    mount({
+      template: `<div></div>`,
+      setup() {
+        useWorkflowAutoSave()
+        return {}
+      }
+    })
+
+    const serviceInstance = (useWorkflowService as any).mock.results[0].value
+    let resolveSave: () => void
+    const firstSavePromise = new Promise<void>((resolve) => {
+      resolveSave = resolve
+    })
+
+    serviceInstance.saveWorkflow.mockImplementationOnce(() => firstSavePromise)
+
+    vi.advanceTimersByTime(1000)
+
+    const graphChangedCallback = (api.addEventListener as any).mock.calls[0][1]
+    graphChangedCallback()
+
+    resolveSave!()
+    await Promise.resolve()
+
+    vi.advanceTimersByTime(1000)
+    expect(serviceInstance.saveWorkflow).toHaveBeenCalledTimes(2)
+  })
+
+  it('should clean up event listeners on component unmount', async () => {
+    mockAutoSaveSetting = 'after delay'
+
+    const wrapper = mount({
+      template: `<div></div>`,
+      setup() {
+        useWorkflowAutoSave()
+        return {}
+      }
+    })
+
+    wrapper.unmount()
+
+    expect(api.removeEventListener).toHaveBeenCalled()
+  })
+
+  it('should handle edge case delay values properly', async () => {
+    mockAutoSaveSetting = 'after delay'
+    mockAutoSaveDelay = 0
+    mockActiveWorkflow = { isModified: true }
+
+    mount({
+      template: `<div></div>`,
+      setup() {
+        useWorkflowAutoSave()
+        return {}
+      }
+    })
+
+    await vi.runAllTimersAsync()
+
+    const serviceInstance = (useWorkflowService as any).mock.results[0].value
+    expect(serviceInstance.saveWorkflow).toHaveBeenCalledTimes(1)
+    serviceInstance.saveWorkflow.mockClear()
+
+    mockAutoSaveDelay = -500
+
+    const graphChangedCallback = (api.addEventListener as any).mock.calls[0][1]
+    graphChangedCallback()
+
+    await vi.runAllTimersAsync()
+
+    expect(serviceInstance.saveWorkflow).toHaveBeenCalledTimes(1)
+  })
+})