[WIP] Persist all unsaved workflow tabs (make tabs behave like vscode)

src/platform/workflow/management/stores/workflowStore.ts

@@ -8,6 +8,7 @@ import type {
   LGraphNode,
   Subgraph
 } from '@/lib/litegraph/src/litegraph'
+import { useWorkflowDraftStore } from '@/platform/workflow/persistence/stores/workflowDraftStore'
 import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/workflowSchema'
 import type { NodeId } from '@/platform/workflow/validation/schemas/workflowSchema'
 import { useWorkflowThumbnail } from '@/renderer/core/thumbnail/useWorkflowThumbnail'
@@ -83,20 +84,45 @@ export class ComfyWorkflow extends UserFile {
   override async load({
     force = false
   }: { force?: boolean } = {}): Promise<LoadedComfyWorkflow> {
+    const draftStore = useWorkflowDraftStore()
+    let draft = !force ? draftStore.getDraft(this.path) : undefined
+    let draftState: ComfyWorkflowJSON | null = null
+    let draftContent: string | null = null
+
+    if (draft) {
+      if (draft.updatedAt < this.lastModified) {
+        draftStore.removeDraft(this.path)
+        draft = undefined
+      }
+    }
+
+    if (draft) {
+      try {
+        draftState = JSON.parse(draft.data)
+        draftContent = draft.data
+      } catch (err) {
+        console.warn('Failed to parse workflow draft, clearing it', err)
+        draftStore.removeDraft(this.path)
+      }
+    }
+
     await super.load({ force })
     if (!force && this.isLoaded) return this as LoadedComfyWorkflow
 
     if (!this.originalContent) {
       throw new Error('[ASSERT] Workflow content should be loaded')
     }
 
-    // Note: originalContent is populated by super.load()
-    this.changeTracker = markRaw(
-      new ChangeTracker(
-        this,
-        /* initialState= */ JSON.parse(this.originalContent)
-      )
-    )
+    const initialState = JSON.parse(this.originalContent)
+    this.changeTracker = markRaw(new ChangeTracker(this, initialState))
+
+    if (draftState && draftContent) {
+      this.changeTracker.activeState = draftState
+      this.content = draftContent
+      this._isModified = true
+      draftStore.markDraftUsed(this.path)
+    }
+
     return this as LoadedComfyWorkflow
   }
 
@@ -106,12 +132,14 @@ export class ComfyWorkflow extends UserFile {
   }
 
   override async save() {
+    const draftStore = useWorkflowDraftStore()
     this.content = JSON.stringify(this.activeState)
     // Force save to ensure the content is updated in remote storage incase
     // the isModified state is screwed by changeTracker.
     const ret = await super.save({ force: true })
     this.changeTracker?.reset()
     this.isModified = false
+    draftStore.removeDraft(this.path)
     return ret
   }
 
@@ -121,8 +149,11 @@ export class ComfyWorkflow extends UserFile {
    * @returns this
    */
   override async saveAs(path: string) {
+    const draftStore = useWorkflowDraftStore()
     this.content = JSON.stringify(this.activeState)
-    return await super.saveAs(path)
+    const result = await super.saveAs(path)
+    draftStore.removeDraft(path)
+    return result
   }
 
   async promptSave(): Promise<string | null> {
@@ -448,6 +479,7 @@ export const useWorkflowStore = defineStore('workflow', () => {
       const oldPath = workflow.path
       const oldKey = workflow.key
       const wasBookmarked = bookmarkStore.isBookmarked(oldPath)
+      const draftStore = useWorkflowDraftStore()
 
       const openIndex = detachWorkflow(workflow)
       // Perform the actual rename operation first
@@ -457,6 +489,8 @@ export const useWorkflowStore = defineStore('workflow', () => {
         attachWorkflow(workflow, openIndex)
       }
 
+      draftStore.moveDraft(oldPath, newPath, workflow.key)
+
       // Move thumbnail from old key to new key (using workflow keys, not full paths)
       const newKey = workflow.key
       moveWorkflowThumbnail(oldKey, newKey)
@@ -474,6 +508,7 @@ export const useWorkflowStore = defineStore('workflow', () => {
     isBusy.value = true
     try {
       await workflow.delete()
+      useWorkflowDraftStore().removeDraft(workflow.path)
       if (bookmarkStore.isBookmarked(workflow.path)) {
         await bookmarkStore.setBookmarked(workflow.path, false)
       }

src/platform/workflow/persistence/base/draftCache.ts

@@ -0,0 +1,77 @@
+export interface WorkflowDraftSnapshot {
+  data: string
+  updatedAt: number
+  name: string
+  isTemporary: boolean
+}
+
+export interface DraftCacheState {
+  drafts: Record<string, WorkflowDraftSnapshot>
+  order: string[]
+}
+
+export const MAX_DRAFTS = 32
+
+export const createDraftCacheState = (
+  drafts: Record<string, WorkflowDraftSnapshot> = {},
+  order: string[] = []
+): DraftCacheState => ({ drafts, order })
+
+export const touchEntry = (order: string[], path: string): string[] => {
+  const next = order.filter((entry) => entry !== path)
+  next.push(path)
+  return next
+}
+
+export const upsertDraft = (
+  state: DraftCacheState,
+  path: string,
+  snapshot: WorkflowDraftSnapshot,
+  limit: number = MAX_DRAFTS
+): DraftCacheState => {
+  const drafts = { ...state.drafts, [path]: snapshot }
+  const order = touchEntry(state.order, path)
+
+  while (order.length > limit) {
+    const oldest = order.shift()
+    if (!oldest) continue
+    if (oldest !== path) {
+      delete drafts[oldest]
+    }
+  }
+
+  return createDraftCacheState(drafts, order)
+}
+
+export const removeDraft = (
+  state: DraftCacheState,
+  path: string
+): DraftCacheState => {
+  if (!(path in state.drafts)) return state
+  const drafts = { ...state.drafts }
+  delete drafts[path]
+  const order = state.order.filter((entry) => entry !== path)
+  return createDraftCacheState(drafts, order)
+}
+
+export const moveDraft = (
+  state: DraftCacheState,
+  oldPath: string,
+  newPath: string,
+  name: string
+): DraftCacheState => {
+  const draft = state.drafts[oldPath]
+  if (!draft) return state
+  const updatedDraft = { ...draft, name }
+  const drafts = { ...state.drafts }
+  delete drafts[oldPath]
+  drafts[newPath] = updatedDraft
+  const order = touchEntry(
+    state.order.filter((entry) => entry !== oldPath && entry !== newPath),
+    newPath
+  )
+  return createDraftCacheState(drafts, order)
+}
+
+export const mostRecentDraftPath = (order: string[]): string | null =>
+  order.length ? order[order.length - 1] : null

src/platform/workflow/persistence/composables/useWorkflowPersistence.ts

@@ -3,7 +3,11 @@ import { computed, watch } from 'vue'
 
 import { useSettingStore } from '@/platform/settings/settingStore'
 import { useWorkflowService } from '@/platform/workflow/core/services/workflowService'
-import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
+import {
+  ComfyWorkflow,
+  useWorkflowStore
+} from '@/platform/workflow/management/stores/workflowStore'
+import { useWorkflowDraftStore } from '@/platform/workflow/persistence/stores/workflowDraftStore'
 import { api } from '@/scripts/api'
 import { app as comfyApp } from '@/scripts/app'
 import { getStorageValue, setStorageValue } from '@/scripts/utils'
@@ -12,45 +16,48 @@ import { useCommandStore } from '@/stores/commandStore'
 export function useWorkflowPersistence() {
   const workflowStore = useWorkflowStore()
   const settingStore = useSettingStore()
+  const workflowDraftStore = useWorkflowDraftStore()
 
   const workflowPersistenceEnabled = computed(() =>
     settingStore.get('Comfy.Workflow.Persist')
   )
 
   const persistCurrentWorkflow = () => {
     if (!workflowPersistenceEnabled.value) return
-    const workflow = JSON.stringify(comfyApp.graph.serialize())
-    localStorage.setItem('workflow', workflow)
+    const activeWorkflow = workflowStore.activeWorkflow
+    if (!activeWorkflow) return
+
+    const graphData = comfyApp.graph.serialize()
+    const workflowJson = JSON.stringify(graphData)
+    localStorage.setItem('workflow', workflowJson)
     if (api.clientId) {
-      sessionStorage.setItem(`workflow:${api.clientId}`, workflow)
+      sessionStorage.setItem(`workflow:${api.clientId}`, workflowJson)
+    }
+
+    if (!activeWorkflow.isTemporary && !activeWorkflow.isModified) {
+      workflowDraftStore.removeDraft(activeWorkflow.path)
+      return
     }
-  }
 
-  const loadWorkflowFromStorage = async (
-    json: string | null,
-    workflowName: string | null
-  ) => {
-    if (!json) return false
-    const workflow = JSON.parse(json)
-    await comfyApp.loadGraphData(workflow, true, true, workflowName)
-    return true
+    workflowDraftStore.saveDraft(activeWorkflow.path, {
+      data: workflowJson,
+      updatedAt: Date.now(),
+      name: activeWorkflow.key,
+      isTemporary: activeWorkflow.isTemporary
+    })
   }
 
   const loadPreviousWorkflowFromStorage = async () => {
     const workflowName = getStorageValue('Comfy.PreviousWorkflow')
-    const clientId = api.initialClientId ?? api.clientId
-
-    // Try loading from session storage first
-    if (clientId) {
-      const sessionWorkflow = sessionStorage.getItem(`workflow:${clientId}`)
-      if (await loadWorkflowFromStorage(sessionWorkflow, workflowName)) {
-        return true
-      }
-    }
-
-    // Fall back to local storage
-    const localWorkflow = localStorage.getItem('workflow')
-    return await loadWorkflowFromStorage(localWorkflow, workflowName)
+    const preferredPath = workflowName
+      ? `${ComfyWorkflow.basePath}${workflowName}`
+      : null
+
+    return await workflowDraftStore.loadPersistedWorkflow({
+      workflowName,
+      preferredPath,
+      fallbackToLatestDraft: !workflowName
+    })
   }
 
   const loadDefaultWorkflow = async () => {
@@ -104,11 +111,12 @@ export function useWorkflowPersistence() {
       }
 
       const paths = openWorkflows.value
-        .filter((workflow) => workflow?.isPersisted)
-        .map((workflow) => workflow.path)
-      const activeIndex = openWorkflows.value.findIndex(
-        (workflow) => workflow.path === activeWorkflow.value?.path
-      )
+        .map((workflow) => workflow?.path)
+        .filter(
+          (path): path is string =>
+            typeof path === 'string' && path.startsWith(ComfyWorkflow.basePath)
+        )
+      const activeIndex = paths.indexOf(activeWorkflow.value.path)
 
       return { paths, activeIndex }
     }
@@ -117,10 +125,10 @@ export function useWorkflowPersistence() {
   // Get storage values before setting watchers
   const storedWorkflows = JSON.parse(
     getStorageValue('Comfy.OpenWorkflowsPaths') || '[]'
-  )
+  ) as string[]
   const storedActiveIndex = JSON.parse(
     getStorageValue('Comfy.ActiveWorkflowIndex') || '-1'
-  )
+  ) as number
 
   watch(restoreState, ({ paths, activeIndex }) => {
     if (workflowPersistenceEnabled.value) {
@@ -132,12 +140,19 @@ export function useWorkflowPersistence() {
   const restoreWorkflowTabsState = () => {
     if (!workflowPersistenceEnabled.value) return
     const isRestorable = storedWorkflows?.length > 0 && storedActiveIndex >= 0
-    if (isRestorable) {
-      workflowStore.openWorkflowsInBackground({
-        left: storedWorkflows.slice(0, storedActiveIndex),
-        right: storedWorkflows.slice(storedActiveIndex)
-      })
-    }
+    if (!isRestorable) return
+
+    storedWorkflows.forEach((path: string) => {
+      if (workflowStore.getWorkflowByPath(path)) return
+      const draft = workflowDraftStore.getDraft(path)
+      if (!draft?.isTemporary) return
+      workflowStore.createTemporary(draft.name)
+    })
+
+    workflowStore.openWorkflowsInBackground({
+      left: storedWorkflows.slice(0, storedActiveIndex),
+      right: storedWorkflows.slice(storedActiveIndex)
+    })
   }
 
   return {