fix: keep node preview images across workflow tab switches

src/platform/workflow/core/services/workflowService.test.ts

@@ -14,6 +14,7 @@ import { useWorkflowStore } from '@/platform/workflow/management/stores/workflow
 import { useWorkflowService } from '@/platform/workflow/core/services/workflowService'
 import { useMissingNodesErrorStore } from '@/platform/nodeReplacement/missingNodesErrorStore'
 import { useExecutionErrorStore } from '@/stores/executionErrorStore'
+import { useNodeOutputStore } from '@/stores/nodeOutputStore'
 import { useMissingModelStore } from '@/platform/missingModel/missingModelStore'
 import { useMissingMediaStore } from '@/platform/missingMedia/missingMediaStore'
 import { app } from '@/scripts/app'
@@ -84,7 +85,9 @@ vi.mock('@/scripts/app', () => ({
   app: {
     canvas: { ds: { offset: [0, 0], scale: 1 } },
     rootGraph: { serialize: vi.fn(() => ({})), extra: {} },
-    loadGraphData: vi.fn()
+    loadGraphData: vi.fn(),
+    nodeOutputs: {},
+    nodePreviewImages: {}
   }
 }))
 
@@ -314,6 +317,21 @@ describe('useWorkflowService', () => {
       )
     })
 
+    it('should stash node previews before app.clean() can revoke them', () => {
+      const activeWorkflow = createModeTestWorkflow({
+        path: 'workflows/test.json'
+      })
+      workflowStore.activeWorkflow = activeWorkflow
+      const stashPreviews = vi.spyOn(
+        useNodeOutputStore(),
+        'stashPreviewsForWorkflow'
+      )
+
+      useWorkflowService().beforeLoadNewGraph()
+
+      expect(stashPreviews).toHaveBeenCalledWith(activeWorkflow.path)
+    })
+
     it('should save active workflow state through the V2 draft store', () => {
       vi.spyOn(useSettingStore(), 'get').mockImplementation((key: string) => {
         return key === 'Comfy.Workflow.Persist'
@@ -561,6 +579,16 @@ describe('useWorkflowService', () => {
       vi.mocked(workflowStore.openWorkflow).mockResolvedValue(existingWorkflow)
     })
 
+    it('should restore the stashed previews of the newly active workflow', async () => {
+      workflowStore.activeWorkflow = existingWorkflow
+
+      await useWorkflowService().afterLoadNewGraph('repeat', makeWorkflowData())
+
+      expect(
+        useNodeOutputStore().restorePreviewsForWorkflow
+      ).toHaveBeenCalledWith(existingWorkflow.path)
+    })
+
     it('should reuse the active workflow when loading the same path repeatedly', async () => {
       const workflowId = 'repeat-workflow-id'
       existingWorkflow.changeTracker.activeState.id = workflowId

src/platform/workflow/core/services/workflowService.ts

@@ -26,6 +26,7 @@ import { useAppMode } from '@/composables/useAppMode'
 import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { useAppModeStore } from '@/stores/appModeStore'
 import { useExecutionErrorStore } from '@/stores/executionErrorStore'
+import { useNodeOutputStore } from '@/stores/nodeOutputStore'
 import { useSubgraphNavigationStore } from '@/stores/subgraphNavigationStore'
 import { useMissingNodesErrorStore } from '@/platform/nodeReplacement/missingNodesErrorStore'
 import { useMissingModelStore } from '@/platform/missingModel/missingModelStore'
@@ -423,6 +424,9 @@ export const useWorkflowService = () => {
         missingMediaCandidates: mediaCandidates ?? undefined
       })
 
+      // Hand the previews to the store before `app.clean()` revokes them
+      useNodeOutputStore().stashPreviewsForWorkflow(activeWorkflow.path)
+
       // Capture thumbnail before loading new graph
       void workflowThumbnail.storeThumbnail(activeWorkflow)
       domWidgetStore.clear()
@@ -447,6 +451,17 @@ export const useWorkflowService = () => {
     value: string | ComfyWorkflow | null,
     workflowData: ComfyWorkflowJSON,
     shareId?: string
+  ) => {
+    await activateLoadedWorkflow(value, workflowData, shareId)
+    useNodeOutputStore().restorePreviewsForWorkflow(
+      useWorkspaceStore().workflow.activeWorkflow?.path
+    )
+  }
+
+  const activateLoadedWorkflow = async (
+    value: string | ComfyWorkflow | null,
+    workflowData: ComfyWorkflowJSON,
+    shareId?: string
   ) => {
     const workflowStore = useWorkspaceStore().workflow
     const { isAppMode } = useAppMode()

src/scripts/app.test.ts

@@ -85,7 +85,9 @@ const {
   },
   mockNodeOutputStore: {
     refreshNodeOutputs: vi.fn(),
-    resetAllOutputsAndPreviews: vi.fn()
+    resetAllOutputsAndPreviews: vi.fn(),
+    stashPreviewsForWorkflow: vi.fn(),
+    restorePreviewsForWorkflow: vi.fn()
   },
   mockWorkspaceWorkflow: {
     activeWorkflow: null as ComfyWorkflow | null,

src/stores/nodeOutputStore.ts

@@ -50,9 +50,15 @@ interface SetOutputOptions {
 }
 
 export const useNodeOutputStore = defineStore('nodeOutput', () => {
-  const { nodeIdToNodeLocatorId, nodeToNodeLocatorId } = useWorkflowStore()
+  const workflowStore = useWorkflowStore()
+  const { nodeIdToNodeLocatorId, nodeToNodeLocatorId } = workflowStore
   const scheduledRevoke: Record<NodeLocatorId, { stop: () => void }> = {}
   const latestPreview = ref<string[]>([])
+  /**
+   * Previews belonging to open-but-inactive workflows, keyed by workflow path.
+   * The stash owns the object URL retain taken when the previews were set.
+   */
+  const stashedPreviews = new Map<string, Record<string, string[]>>()
 
   function scheduleRevoke(locator: NodeLocatorId, cb: () => void) {
     scheduledRevoke[locator]?.stop()
@@ -298,19 +304,74 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     delete nodePreviewImages.value[nodeLocatorId]
   }
 
-  function revokeAllPreviews() {
-    for (const nodeLocatorId of Object.keys(app.nodePreviewImages)) {
-      const previews = app.nodePreviewImages[nodeLocatorId]
-      if (!previews?.[Symbol.iterator]) continue
+  function releasePreviewUrls(previews: Record<string, string[]>) {
+    for (const urls of Object.values(previews)) {
+      if (!urls?.[Symbol.iterator]) continue
 
-      for (const url of previews) {
+      for (const url of urls) {
         releaseSharedObjectUrl(url)
       }
     }
+  }
+
+  function revokeAllPreviews() {
+    releasePreviewUrls(app.nodePreviewImages)
+    app.nodePreviewImages = {}
+    nodePreviewImages.value = {}
+  }
+
+  function discardStashedPreviews(workflowPath: string) {
+    const stashed = stashedPreviews.get(workflowPath)
+    if (!stashed) return
+
+    stashedPreviews.delete(workflowPath)
+    releasePreviewUrls(stashed)
+  }
+
+  function discardClosedWorkflowPreviews() {
+    const openPaths = new Set(workflowStore.openWorkflows.map((wf) => wf.path))
+    for (const path of [...stashedPreviews.keys()]) {
+      if (!openPaths.has(path)) discardStashedPreviews(path)
+    }
+  }
+
+  /**
+   * Move the live previews into the stash so `app.clean()` cannot revoke them
+   * while `workflowPath` is loaded out of the editor. Preview frames arrive
+   * over the websocket and cannot be re-fetched, so releasing their object
+   * URLs on a tab switch loses them for good.
+   */
+  function stashPreviewsForWorkflow(workflowPath: string) {
+    discardClosedWorkflowPreviews()
+    discardStashedPreviews(workflowPath)
+
+    const previews = app.nodePreviewImages
+    if (Object.keys(previews).length) {
+      stashedPreviews.set(workflowPath, previews)
+    }
     app.nodePreviewImages = {}
     nodePreviewImages.value = {}
   }
 
+  function restorePreviewsForWorkflow(workflowPath: string | undefined) {
+    if (!workflowPath) return
+
+    const stashed = stashedPreviews.get(workflowPath)
+    if (!stashed) return
+    stashedPreviews.delete(workflowPath)
+
+    const live = app.nodePreviewImages
+    const superseded: Record<string, string[]> = {}
+    for (const [nodeLocatorId, urls] of Object.entries(stashed)) {
+      // A preview that arrived while the graph was loading wins; the stashed
+      // copy for that node then has no owner left to release it.
+      if (nodeLocatorId in live) superseded[nodeLocatorId] = urls
+      else live[nodeLocatorId] = urls
+    }
+    releasePreviewUrls(superseded)
+    nodePreviewImages.value = { ...live }
+  }
+
   function revokeSubgraphPreviews(subgraphNode: SubgraphNode) {
     const { graph } = subgraphNode
     if (!graph) return
@@ -444,6 +505,8 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     revokePreviewsByLocatorId,
     revokeAllPreviews,
     revokeSubgraphPreviews,
+    stashPreviewsForWorkflow,
+    restorePreviewsForWorkflow,
     removeNodeOutputs,
     removeNodeOutputsForNode,
     snapshotOutputs,

src/stores/nodeOutputStore.workflowSwitch.test.ts

@@ -0,0 +1,164 @@
+import { fromAny } from '@total-typescript/shoehorn'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+import type { LGraphNode } from '@/lib/litegraph/src/litegraph'
+import { app } from '@/scripts/app'
+import { useNodeOutputStore } from '@/stores/nodeOutputStore'
+import { toNodeId } from '@/types/nodeId'
+
+const { WORKFLOW_A, WORKFLOW_B } = vi.hoisted(() => ({
+  WORKFLOW_A: 'workflows/a.json',
+  WORKFLOW_B: 'workflows/b.json'
+}))
+
+const mocks = vi.hoisted(() => ({
+  workflowStore: null as unknown as {
+    activeWorkflow: { path: string } | null
+    openWorkflows: { path: string }[]
+    nodeIdToNodeLocatorId: (id: string | number) => string
+    nodeToNodeLocatorId: (node: { id: string | number }) => string
+  }
+}))
+
+vi.mock('@/utils/litegraphUtil', () => ({
+  isAnimatedOutput: vi.fn(() => false),
+  isVideoNode: vi.fn(() => false),
+  resolveNode: vi.fn()
+}))
+
+vi.mock('@/utils/graphTraversalUtil', () => ({
+  executionIdToNodeLocatorId: vi.fn((_rootGraph: unknown, id: string) => id)
+}))
+
+vi.mock('@/scripts/app', () => ({
+  app: {
+    getPreviewFormatParam: vi.fn(() => ''),
+    getRandParam: vi.fn(() => ''),
+    rootGraph: { getNodeById: vi.fn() },
+    nodeOutputs: {} as Record<string, unknown>,
+    nodePreviewImages: {} as Record<string, string[]>
+  }
+}))
+
+vi.mock('@/platform/workflow/management/stores/workflowStore', async () => {
+  const { reactive } = await import('vue')
+  mocks.workflowStore = reactive({
+    activeWorkflow: { path: WORKFLOW_A } as { path: string } | null,
+    openWorkflows: [{ path: WORKFLOW_A }, { path: WORKFLOW_B }],
+    nodeIdToNodeLocatorId: (id: string | number) => String(id),
+    nodeToNodeLocatorId: (node: { id: string | number }) => String(node.id)
+  })
+  return { useWorkflowStore: () => mocks.workflowStore }
+})
+
+const createMockNode = (id: number): LGraphNode =>
+  fromAny<LGraphNode, unknown>({ id: toNodeId(id), type: 'KSampler' })
+
+/**
+ * Mirrors `workflowService.beforeLoadNewGraph()` -> `app.clean()` ->
+ * `workflowService.afterLoadNewGraph()`, the sequence every workflow load and
+ * tab switch runs.
+ */
+function switchToWorkflow(path: string) {
+  const store = useNodeOutputStore()
+  const leaving = mocks.workflowStore.activeWorkflow
+  if (leaving) store.stashPreviewsForWorkflow(leaving.path)
+  store.resetAllOutputsAndPreviews()
+  mocks.workflowStore.activeWorkflow = { path }
+  store.restorePreviewsForWorkflow(path)
+}
+
+// Object URL retain counts live in a module-level map in objectUrlUtil, so
+// every test needs its own URLs or the counts leak between them.
+let urlCounter = 0
+
+describe('nodeOutputStore preview lifecycle across workflow tab switches', () => {
+  let revokeObjectURL: ReturnType<typeof vi.spyOn>
+
+  const createBlobUrl = () => URL.createObjectURL(new Blob(['x']))
+
+  beforeEach(() => {
+    vi.spyOn(URL, 'createObjectURL').mockImplementation(
+      () => `blob:mock/${++urlCounter}`
+    )
+    revokeObjectURL = vi
+      .spyOn(URL, 'revokeObjectURL')
+      .mockImplementation(() => {})
+    app.nodeOutputs = {}
+    app.nodePreviewImages = {}
+    mocks.workflowStore.activeWorkflow = { path: WORKFLOW_A }
+    mocks.workflowStore.openWorkflows = [
+      { path: WORKFLOW_A },
+      { path: WORKFLOW_B }
+    ]
+  })
+
+  it('keeps a finished run preview when the user switches tabs and comes back', () => {
+    const store = useNodeOutputStore()
+    const node = createMockNode(5)
+    const previewUrl = createBlobUrl()
+
+    // A run finished on workflow A and left its last preview frame on node 5.
+    // Nothing revokes it at execution end, so the node keeps displaying it.
+    store.setNodePreviewsByNodeId(node.id, [previewUrl])
+    expect(store.getNodePreviews(node)).toEqual([previewUrl])
+
+    // User switches to workflow B, then back to workflow A.
+    switchToWorkflow(WORKFLOW_B)
+    switchToWorkflow(WORKFLOW_A)
+
+    expect(store.getNodePreviews(node)).toEqual([previewUrl])
+    expect(store.nodePreviewImages['5']).toEqual([previewUrl])
+    // The blob must still be alive; a b_preview frame cannot be re-fetched.
+    expect(revokeObjectURL).not.toHaveBeenCalledWith(previewUrl)
+  })
+
+  it('does not revoke the preview blob when leaving the workflow', () => {
+    const store = useNodeOutputStore()
+    const previewUrl = createBlobUrl()
+
+    store.setNodePreviewsByNodeId(createMockNode(5).id, [previewUrl])
+    switchToWorkflow(WORKFLOW_B)
+
+    // Revoking makes the loss permanent: there is no source to re-derive a
+    // websocket preview frame from once the object URL is gone.
+    expect(revokeObjectURL).not.toHaveBeenCalledWith(previewUrl)
+  })
+
+  it('does not show one workflow preview on another workflow same-id node', () => {
+    const store = useNodeOutputStore()
+    const node = createMockNode(5)
+
+    store.setNodePreviewsByNodeId(node.id, [createBlobUrl()])
+    switchToWorkflow(WORKFLOW_B)
+
+    expect(store.getNodePreviews(node)).toBeUndefined()
+    expect(store.nodePreviewImages).toEqual({})
+  })
+
+  it('releases stashed previews once the workflow is closed', () => {
+    const store = useNodeOutputStore()
+    const previewUrl = createBlobUrl()
+
+    store.setNodePreviewsByNodeId(createMockNode(5).id, [previewUrl])
+    switchToWorkflow(WORKFLOW_B)
+
+    mocks.workflowStore.openWorkflows = [{ path: WORKFLOW_B }]
+    switchToWorkflow(WORKFLOW_B)
+
+    expect(revokeObjectURL).toHaveBeenCalledWith(previewUrl)
+  })
+
+  it('still revokes previews when the workflow is cleared in place', () => {
+    const store = useNodeOutputStore()
+    const node = createMockNode(5)
+    const previewUrl = createBlobUrl()
+
+    store.setNodePreviewsByNodeId(node.id, [previewUrl])
+    // "Clear Workflow" calls app.clean() without loading another graph.
+    store.resetAllOutputsAndPreviews()
+
+    expect(store.getNodePreviews(node)).toBeUndefined()
+    expect(revokeObjectURL).toHaveBeenCalledWith(previewUrl)
+  })
+})