Skip to content

Commit 4c28897

Browse files
fix: keep node preview images across workflow tab switches (#15360)
## Summary Node preview images are lost permanently on a workflow tab switch, because `app.clean()` revokes their object URLs and nothing restores them. This scopes preview state per workflow instead. Fixes FE-1645 ## Changes - **What**: `workflowService.beforeLoadNewGraph()` hands the live previews to `nodeOutputStore` keyed by the outgoing workflow path, so `app.clean()` has nothing left to revoke. `afterLoadNewGraph()` puts back the previews belonging to the workflow that just became active. ### Before 1. Open workflow A with a KSampler, preview method anything but `none`. 2. Queue it and let it finish — the KSampler shows its final latent preview. 3. Switch to tab B, switch back to tab A. 4. The preview is gone and never comes back. Output images on other nodes survive the same switch. `app.clean()` (`src/scripts/app.ts:2434`) runs on every workflow load and calls `resetAllOutputsAndPreviews()` → `revokeAllPreviews()` (`src/stores/nodeOutputStore.ts:301`), which releases every preview object URL and empties the maps. `nodeOutputs` survives only because `ChangeTracker` snapshots and restores it (`src/scripts/changeTracker.ts:306`/`350`); previews have no equivalent. Nothing revokes previews at execution end, so a finished run's last frame is real state a user is looking at — and a `b_preview` websocket frame cannot be re-fetched, so revoking makes the loss permanent. ### After The preview is still there when you come back to the tab. ## Review Focus - Keying by workflow path is load-bearing, not incidental. Node locator ids for root-graph nodes are bare node ids, so simply not clearing on load would show workflow A's node 5 preview on workflow B's node 5. - Ownership of the object URL retain moves from the live map to the stash and back, so nothing is released on a tab switch. Released instead when: the workflow is cleared in place (`Clear Workflow` still calls `app.clean()` with no load following it), the workflow has been closed, or a preview arriving for the incoming graph supersedes the stashed one. - Undo/redo goes through `loadGraphData(..., clean = false)` for the same workflow — it stashes and restores under the same path, so the net effect is unchanged. - Deliberately not done: no `ChangeTracker` snapshot/restore field for previews, and no use of the `workflowTransientState` mechanism proposed in #14941, so this does not prejudge that review. ## Tests `src/stores/nodeOutputStore.workflowSwitch.test.ts` covers the round trip, the same-node-id cross-workflow case, release on close, and clear-in-place. `workflowService.test.ts` covers the two lifecycle hooks.
1 parent 49e3620 commit 4c28897

5 files changed

Lines changed: 344 additions & 8 deletions

File tree

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

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import { useWorkflowStore } from '@/platform/workflow/management/stores/workflow
1414
import { useWorkflowService } from '@/platform/workflow/core/services/workflowService'
1515
import { useMissingNodesErrorStore } from '@/platform/nodeReplacement/missingNodesErrorStore'
1616
import { useExecutionErrorStore } from '@/stores/executionErrorStore'
17+
import { useNodeOutputStore } from '@/stores/nodeOutputStore'
1718
import { useMissingModelStore } from '@/platform/missingModel/missingModelStore'
1819
import { useMissingMediaStore } from '@/platform/missingMedia/missingMediaStore'
1920
import { app } from '@/scripts/app'
@@ -84,7 +85,9 @@ vi.mock('@/scripts/app', () => ({
8485
app: {
8586
canvas: { ds: { offset: [0, 0], scale: 1 } },
8687
rootGraph: { serialize: vi.fn(() => ({})), extra: {} },
87-
loadGraphData: vi.fn()
88+
loadGraphData: vi.fn(),
89+
nodeOutputs: {},
90+
nodePreviewImages: {}
8891
}
8992
}))
9093

@@ -314,6 +317,21 @@ describe('useWorkflowService', () => {
314317
)
315318
})
316319

320+
it('should stash node previews before app.clean() can revoke them', () => {
321+
const activeWorkflow = createModeTestWorkflow({
322+
path: 'workflows/test.json'
323+
})
324+
workflowStore.activeWorkflow = activeWorkflow
325+
const stashPreviews = vi.spyOn(
326+
useNodeOutputStore(),
327+
'stashPreviewsForWorkflow'
328+
)
329+
330+
useWorkflowService().beforeLoadNewGraph()
331+
332+
expect(stashPreviews).toHaveBeenCalledWith(activeWorkflow.path)
333+
})
334+
317335
it('should save active workflow state through the V2 draft store', () => {
318336
vi.spyOn(useSettingStore(), 'get').mockImplementation((key: string) => {
319337
return key === 'Comfy.Workflow.Persist'
@@ -542,6 +560,20 @@ describe('useWorkflowService', () => {
542560
expect(closed).toBe(false)
543561
expect(workflowStore.closeWorkflow).not.toHaveBeenCalled()
544562
})
563+
564+
it('should release the closed workflow node previews', async () => {
565+
const workflow = createModeTestWorkflow({
566+
path: 'workflows/closing.json'
567+
})
568+
const discardPreviews = vi.spyOn(
569+
useNodeOutputStore(),
570+
'discardPreviewsForWorkflow'
571+
)
572+
573+
await service.closeWorkflow(workflow, { warnIfUnsaved: false })
574+
575+
expect(discardPreviews).toHaveBeenCalledWith(workflow.path)
576+
})
545577
})
546578

547579
describe('afterLoadNewGraph', () => {
@@ -561,6 +593,16 @@ describe('useWorkflowService', () => {
561593
vi.mocked(workflowStore.openWorkflow).mockResolvedValue(existingWorkflow)
562594
})
563595

596+
it('should restore the stashed previews of the newly active workflow', async () => {
597+
workflowStore.activeWorkflow = existingWorkflow
598+
599+
await useWorkflowService().afterLoadNewGraph('repeat', makeWorkflowData())
600+
601+
expect(
602+
useNodeOutputStore().restorePreviewsForWorkflow
603+
).toHaveBeenCalledWith(existingWorkflow.path)
604+
})
605+
564606
it('should reuse the active workflow when loading the same path repeatedly', async () => {
565607
const workflowId = 'repeat-workflow-id'
566608
existingWorkflow.changeTracker.activeState.id = workflowId

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import { useAppMode } from '@/composables/useAppMode'
2626
import { useDomWidgetStore } from '@/stores/domWidgetStore'
2727
import { useAppModeStore } from '@/stores/appModeStore'
2828
import { useExecutionErrorStore } from '@/stores/executionErrorStore'
29+
import { useNodeOutputStore } from '@/stores/nodeOutputStore'
2930
import { useSubgraphNavigationStore } from '@/stores/subgraphNavigationStore'
3031
import { useMissingNodesErrorStore } from '@/platform/nodeReplacement/missingNodesErrorStore'
3132
import { useMissingModelStore } from '@/platform/missingModel/missingModelStore'
@@ -348,6 +349,7 @@ export const useWorkflowService = () => {
348349
}
349350

350351
await workflowStore.closeWorkflow(workflow)
352+
useNodeOutputStore().discardPreviewsForWorkflow(workflow.path)
351353
return true
352354
}
353355

@@ -423,6 +425,9 @@ export const useWorkflowService = () => {
423425
missingMediaCandidates: mediaCandidates ?? undefined
424426
})
425427

428+
// Hand the previews to the store before `app.clean()` revokes them
429+
useNodeOutputStore().stashPreviewsForWorkflow(activeWorkflow.path)
430+
426431
// Capture thumbnail before loading new graph
427432
void workflowThumbnail.storeThumbnail(activeWorkflow)
428433
domWidgetStore.clear()
@@ -447,6 +452,17 @@ export const useWorkflowService = () => {
447452
value: string | ComfyWorkflow | null,
448453
workflowData: ComfyWorkflowJSON,
449454
shareId?: string
455+
) => {
456+
await activateLoadedWorkflow(value, workflowData, shareId)
457+
useNodeOutputStore().restorePreviewsForWorkflow(
458+
useWorkspaceStore().workflow.activeWorkflow?.path
459+
)
460+
}
461+
462+
const activateLoadedWorkflow = async (
463+
value: string | ComfyWorkflow | null,
464+
workflowData: ComfyWorkflowJSON,
465+
shareId?: string
450466
) => {
451467
const workflowStore = useWorkspaceStore().workflow
452468
const { isAppMode } = useAppMode()

src/scripts/app.test.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,10 @@ const {
8585
},
8686
mockNodeOutputStore: {
8787
refreshNodeOutputs: vi.fn(),
88-
resetAllOutputsAndPreviews: vi.fn()
88+
resetAllOutputsAndPreviews: vi.fn(),
89+
stashPreviewsForWorkflow: vi.fn(),
90+
restorePreviewsForWorkflow: vi.fn(),
91+
discardPreviewsForWorkflow: vi.fn()
8992
},
9093
mockWorkspaceWorkflow: {
9194
activeWorkflow: null as ComfyWorkflow | null,

src/stores/nodeOutputStore.ts

Lines changed: 71 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,15 @@ interface SetOutputOptions {
5050
}
5151

5252
export const useNodeOutputStore = defineStore('nodeOutput', () => {
53-
const { nodeIdToNodeLocatorId, nodeToNodeLocatorId } = useWorkflowStore()
53+
const workflowStore = useWorkflowStore()
54+
const { nodeIdToNodeLocatorId, nodeToNodeLocatorId } = workflowStore
5455
const scheduledRevoke: Record<NodeLocatorId, { stop: () => void }> = {}
5556
const latestPreview = ref<string[]>([])
57+
/**
58+
* Previews belonging to open-but-inactive workflows, keyed by workflow path.
59+
* The stash owns the object URL retain taken when the previews were set.
60+
*/
61+
const stashedPreviews = new Map<string, Record<string, string[]>>()
5662

5763
function scheduleRevoke(locator: NodeLocatorId, cb: () => void) {
5864
scheduledRevoke[locator]?.stop()
@@ -298,19 +304,75 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
298304
delete nodePreviewImages.value[nodeLocatorId]
299305
}
300306

301-
function revokeAllPreviews() {
302-
for (const nodeLocatorId of Object.keys(app.nodePreviewImages)) {
303-
const previews = app.nodePreviewImages[nodeLocatorId]
304-
if (!previews?.[Symbol.iterator]) continue
307+
function releasePreviewUrls(previews: Record<string, string[]>) {
308+
for (const urls of Object.values(previews)) {
309+
if (!urls?.[Symbol.iterator]) continue
305310

306-
for (const url of previews) {
311+
for (const url of urls) {
307312
releaseSharedObjectUrl(url)
308313
}
309314
}
315+
}
316+
317+
function revokeAllPreviews() {
318+
releasePreviewUrls(app.nodePreviewImages)
319+
app.nodePreviewImages = {}
320+
nodePreviewImages.value = {}
321+
}
322+
323+
/** Release the previews held for a workflow that is going away. */
324+
function discardPreviewsForWorkflow(workflowPath: string) {
325+
const stashed = stashedPreviews.get(workflowPath)
326+
if (!stashed) return
327+
328+
stashedPreviews.delete(workflowPath)
329+
releasePreviewUrls(stashed)
330+
}
331+
332+
function discardClosedWorkflowPreviews() {
333+
const openPaths = new Set(workflowStore.openWorkflows.map((wf) => wf.path))
334+
for (const path of [...stashedPreviews.keys()]) {
335+
if (!openPaths.has(path)) discardPreviewsForWorkflow(path)
336+
}
337+
}
338+
339+
/**
340+
* Move the live previews into the stash so `app.clean()` cannot revoke them
341+
* while `workflowPath` is loaded out of the editor. Preview frames arrive
342+
* over the websocket and cannot be re-fetched, so releasing their object
343+
* URLs on a tab switch loses them for good.
344+
*/
345+
function stashPreviewsForWorkflow(workflowPath: string) {
346+
discardClosedWorkflowPreviews()
347+
discardPreviewsForWorkflow(workflowPath)
348+
349+
const previews = app.nodePreviewImages
350+
if (Object.keys(previews).length) {
351+
stashedPreviews.set(workflowPath, previews)
352+
}
310353
app.nodePreviewImages = {}
311354
nodePreviewImages.value = {}
312355
}
313356

357+
function restorePreviewsForWorkflow(workflowPath: string | undefined) {
358+
if (!workflowPath) return
359+
360+
const stashed = stashedPreviews.get(workflowPath)
361+
if (!stashed) return
362+
stashedPreviews.delete(workflowPath)
363+
364+
const live = app.nodePreviewImages
365+
const superseded: Record<string, string[]> = {}
366+
for (const [nodeLocatorId, urls] of Object.entries(stashed)) {
367+
// A preview that arrived while the graph was loading wins; the stashed
368+
// copy for that node then has no owner left to release it.
369+
if (nodeLocatorId in live) superseded[nodeLocatorId] = urls
370+
else live[nodeLocatorId] = urls
371+
}
372+
releasePreviewUrls(superseded)
373+
nodePreviewImages.value = { ...live }
374+
}
375+
314376
function revokeSubgraphPreviews(subgraphNode: SubgraphNode) {
315377
const { graph } = subgraphNode
316378
if (!graph) return
@@ -444,6 +506,9 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
444506
revokePreviewsByLocatorId,
445507
revokeAllPreviews,
446508
revokeSubgraphPreviews,
509+
stashPreviewsForWorkflow,
510+
restorePreviewsForWorkflow,
511+
discardPreviewsForWorkflow,
447512
removeNodeOutputs,
448513
removeNodeOutputsForNode,
449514
snapshotOutputs,

0 commit comments

Comments
 (0)