From f83daa6f3b3a1de31222d515bc7dc4aa28ff4822 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Fri, 27 Feb 2026 18:58:41 +0000 Subject: [PATCH] App mode - discard slow preview messages to prevent overwriting output image (#9261) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Prevent latent previews received after the job/node has already finished processing overwriting the actual output display ## Changes - **What**: - updates job preview store to also track which node the preview was for - updates linear progress tracking to store executed nodes enabling skipping previews of these ## Review Focus ## Screenshots (if applicable) ┆Issue is synchronized with this [Notion page](https://www.notion.so/PR-9261-App-mode-discard-slow-preview-messages-to-prevent-overwriting-output-image-3136d73d3650817884c2ce2ff5993b9e) by [Unito](https://www.unito.io) --- .../linearMode/linearOutputStore.test.ts | 115 ++++++++++++++++-- .../linearMode/linearOutputStore.ts | 25 +++- src/scripts/app.ts | 2 +- src/stores/jobPreviewStore.test.ts | 102 ++++++++++++++++ src/stores/jobPreviewStore.ts | 50 +++++--- 5 files changed, 265 insertions(+), 29 deletions(-) create mode 100644 src/stores/jobPreviewStore.test.ts diff --git a/src/renderer/extensions/linearMode/linearOutputStore.test.ts b/src/renderer/extensions/linearMode/linearOutputStore.test.ts index 0efbb84165..71eafef5b8 100644 --- a/src/renderer/extensions/linearMode/linearOutputStore.test.ts +++ b/src/renderer/extensions/linearMode/linearOutputStore.test.ts @@ -7,7 +7,7 @@ import type { ExecutedWsMessage } from '@/schemas/apiSchema' import { ResultItemImpl } from '@/stores/queueStore' const activeJobIdRef = ref(null) -const previewsRef = ref>({}) +const previewsRef = ref>({}) const isAppModeRef = ref(true) const { apiTarget } = vi.hoisted(() => ({ @@ -30,7 +30,7 @@ vi.mock('@/stores/executionStore', () => ({ vi.mock('@/stores/jobPreviewStore', () => ({ useJobPreviewStore: () => ({ - get previewsByPromptId() { + get nodePreviewsByPromptId() { return previewsRef.value } }) @@ -63,12 +63,13 @@ function makeExecutedDetail( promptId: string, images: Array> = [ { filename: 'out.png', subfolder: '', type: 'output' } - ] + ], + nodeId = '1' ): ExecutedWsMessage { return { prompt_id: promptId, - node: '1', - display_node: '1', + node: nodeId, + display_node: nodeId, output: { images } } as ExecutedWsMessage } @@ -329,7 +330,7 @@ describe('linearOutputStore', () => { expect(store.selectedId).toBeNull() }) - it('transitions to latent via previewsByPromptId watcher', async () => { + it('transitions to latent via previews watcher', async () => { vi.useFakeTimers() const { nextTick } = await import('vue') const store = useLinearOutputStore() @@ -341,7 +342,9 @@ describe('linearOutputStore', () => { expect(store.inProgressItems[0].state).toBe('skeleton') // Simulate jobPreviewStore update - previewsRef.value = { 'job-1': 'blob:preview-1' } + previewsRef.value = { + 'job-1': { url: 'blob:preview-1', nodeId: 'node-1' } + } await nextTick() vi.advanceTimersByTime(16) @@ -487,6 +490,104 @@ describe('linearOutputStore', () => { expect(store.inProgressItems).toHaveLength(0) }) + it('discards latent previews for already-executed nodes', () => { + vi.useFakeTimers() + const store = useLinearOutputStore() + store.onJobStart('job-1') + + // Node 1 sends latent then executes + store.onLatentPreview('job-1', 'blob:node1-latent', '1') + vi.advanceTimersByTime(16) + store.onNodeExecuted('job-1', makeExecutedDetail('job-1', undefined, '1')) + + // Stale latent for node 1 arrives after it already executed + store.onLatentPreview('job-1', 'blob:node1-stale', '1') + vi.advanceTimersByTime(16) + + // Should not create a new latent item for the executed node + expect( + store.inProgressItems.filter((i) => i.state === 'latent') + ).toHaveLength(0) + vi.useRealTimers() + }) + + it('accepts latent previews for new nodes after prior node executed', () => { + vi.useFakeTimers() + const store = useLinearOutputStore() + store.onJobStart('job-1') + + // Node 1 executes + store.onNodeExecuted('job-1', makeExecutedDetail('job-1', undefined, '1')) + + // Node 2 sends latent preview — should be accepted + store.onLatentPreview('job-1', 'blob:node2-latent', '2') + vi.advanceTimersByTime(16) + + expect( + store.inProgressItems.filter((i) => i.state === 'latent') + ).toHaveLength(1) + expect(store.inProgressItems[0].latentPreviewUrl).toBe('blob:node2-latent') + vi.useRealTimers() + }) + + it('cancels pending RAF when a node executes', () => { + vi.useFakeTimers() + const store = useLinearOutputStore() + store.onJobStart('job-1') + + // Latent preview scheduled in RAF + store.onLatentPreview('job-1', 'blob:node1-latent') + // Node executes before RAF fires — should cancel it + store.onNodeExecuted('job-1', makeExecutedDetail('job-1', undefined, '1')) + vi.advanceTimersByTime(16) + + // Only the image item, no latent + expect( + store.inProgressItems.filter((i) => i.state === 'latent') + ).toHaveLength(0) + expect( + store.inProgressItems.filter((i) => i.state === 'image') + ).toHaveLength(1) + vi.useRealTimers() + }) + + it('discards latent previews arriving after job completion', () => { + vi.useFakeTimers() + const store = useLinearOutputStore() + store.onJobStart('job-1') + store.onNodeExecuted('job-1', makeExecutedDetail('job-1')) + + // Latent preview scheduled in RAF before job completes + store.onLatentPreview('job-1', 'blob:late') + store.onJobComplete('job-1') + + // RAF fires after completion — should be cancelled + vi.advanceTimersByTime(16) + + // No new latent items should have been created + expect( + store.inProgressItems.filter((i) => i.state === 'latent') + ).toHaveLength(0) + vi.useRealTimers() + }) + + it('discards latent previews for completed job after RAF', () => { + vi.useFakeTimers() + const store = useLinearOutputStore() + store.onJobStart('job-1') + store.onNodeExecuted('job-1', makeExecutedDetail('job-1')) + store.onJobComplete('job-1') + + // Late preview arrives after job already completed + store.onLatentPreview('job-1', 'blob:very-late') + vi.advanceTimersByTime(16) + + expect( + store.inProgressItems.filter((i) => i.state === 'latent') + ).toHaveLength(0) + vi.useRealTimers() + }) + it('ignores executed events for other jobs', () => { const store = useLinearOutputStore() store.onJobStart('job-1') diff --git a/src/renderer/extensions/linearMode/linearOutputStore.ts b/src/renderer/extensions/linearMode/linearOutputStore.ts index 76382a54b2..9eb6cbee3c 100644 --- a/src/renderer/extensions/linearMode/linearOutputStore.ts +++ b/src/renderer/extensions/linearMode/linearOutputStore.ts @@ -19,6 +19,7 @@ export const useLinearOutputStore = defineStore('linearOutput', () => { const isFollowing = ref(true) const trackedJobId = ref(null) const pendingResolve = ref(new Set()) + const executedNodeIds = new Set() let nextSeq = 0 @@ -40,6 +41,7 @@ export const useLinearOutputStore = defineStore('linearOutput', () => { const currentSkeletonId = shallowRef(null) function onJobStart(jobId: string) { + executedNodeIds.clear() const item: InProgressItem = { id: makeItemId(jobId), jobId, @@ -53,7 +55,9 @@ export const useLinearOutputStore = defineStore('linearOutput', () => { } let raf: number | null = null - function onLatentPreview(jobId: string, url: string) { + function onLatentPreview(jobId: string, url: string, nodeId?: string) { + if (nodeId && executedNodeIds.has(nodeId)) return + // Issue in Firefox where it doesnt seem to always re-render, wrapping in RAF fixes it if (raf) cancelAnimationFrame(raf) raf = requestAnimationFrame(() => { @@ -89,6 +93,11 @@ export const useLinearOutputStore = defineStore('linearOutput', () => { function onNodeExecuted(jobId: string, detail: ExecutedWsMessage) { const nodeId = String(detail.display_node || detail.node) + executedNodeIds.add(nodeId) + if (raf) { + cancelAnimationFrame(raf) + raf = null + } const newOutputs = flattenNodeOutput([nodeId, detail.output]) if (newOutputs.length === 0) return @@ -134,7 +143,14 @@ export const useLinearOutputStore = defineStore('linearOutput', () => { } function onJobComplete(jobId: string) { + if (raf) { + cancelAnimationFrame(raf) + raf = null + } currentSkeletonId.value = null + if (trackedJobId.value === jobId) { + trackedJobId.value = null + } const hasImages = inProgressItems.value.some( (i) => i.jobId === jobId && i.state === 'image' @@ -209,6 +225,7 @@ export const useLinearOutputStore = defineStore('linearOutput', () => { cancelAnimationFrame(raf) raf = null } + executedNodeIds.clear() inProgressItems.value = [] selectedId.value = null isFollowing.value = true @@ -231,13 +248,13 @@ export const useLinearOutputStore = defineStore('linearOutput', () => { ) watch( - () => jobPreviewStore.previewsByPromptId, + () => jobPreviewStore.nodePreviewsByPromptId, (previews) => { if (!isAppMode.value) return const jobId = executionStore.activeJobId if (!jobId) return - const url = previews[jobId] - if (url) onLatentPreview(jobId, url) + const preview = previews[jobId] + if (preview) onLatentPreview(jobId, preview.url, preview.nodeId) }, { deep: true } ) diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 14851afea4..ca12dadac7 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -761,7 +761,7 @@ export class ComfyApp { const { setNodePreviewsByExecutionId, revokePreviewsByExecutionId } = useNodeOutputStore() const blobUrl = createSharedObjectUrl(blob) - useJobPreviewStore().setPreviewUrl(jobId, blobUrl) + useJobPreviewStore().setPreviewUrl(jobId, blobUrl, displayNodeId) // Ensure clean up if `executing` event is missed. revokePreviewsByExecutionId(displayNodeId) // Preview cleanup is handled in progress_state event to support multiple concurrent previews diff --git a/src/stores/jobPreviewStore.test.ts b/src/stores/jobPreviewStore.test.ts new file mode 100644 index 0000000000..b10ff84d04 --- /dev/null +++ b/src/stores/jobPreviewStore.test.ts @@ -0,0 +1,102 @@ +import { createPinia, setActivePinia } from 'pinia' +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { ref } from 'vue' + +import { useJobPreviewStore } from '@/stores/jobPreviewStore' +import { releaseSharedObjectUrl } from '@/utils/objectUrlUtil' + +const previewMethodRef = ref('latent2rgb') + +vi.mock('@/platform/settings/settingStore', () => ({ + useSettingStore: () => ({ + get: (key: string) => { + if (key === 'Comfy.Execution.PreviewMethod') return previewMethodRef.value + return undefined + } + }) +})) + +vi.mock('@/utils/objectUrlUtil', () => ({ + retainSharedObjectUrl: vi.fn(), + releaseSharedObjectUrl: vi.fn() +})) + +describe('jobPreviewStore', () => { + beforeEach(() => { + vi.clearAllMocks() + setActivePinia(createPinia()) + previewMethodRef.value = 'latent2rgb' + }) + + it('stores preview with nodeId', () => { + const store = useJobPreviewStore() + store.setPreviewUrl('prompt-1', 'blob:url-1', 'node-5') + + expect(store.nodePreviewsByPromptId['prompt-1']).toEqual({ + url: 'blob:url-1', + nodeId: 'node-5' + }) + }) + + it('stores preview without nodeId', () => { + const store = useJobPreviewStore() + store.setPreviewUrl('prompt-1', 'blob:url-1') + + expect(store.nodePreviewsByPromptId['prompt-1']).toEqual({ + url: 'blob:url-1', + nodeId: undefined + }) + }) + + it('derives previewsByPromptId as url-only map', () => { + const store = useJobPreviewStore() + store.setPreviewUrl('p1', 'blob:a', 'node-1') + store.setPreviewUrl('p2', 'blob:b', 'node-2') + + expect(store.previewsByPromptId).toEqual({ + p1: 'blob:a', + p2: 'blob:b' + }) + }) + + it('clears a single preview', () => { + const store = useJobPreviewStore() + store.setPreviewUrl('p1', 'blob:a', 'node-1') + store.setPreviewUrl('p2', 'blob:b', 'node-2') + + store.clearPreview('p1') + + expect(store.nodePreviewsByPromptId['p1']).toBeUndefined() + expect(store.nodePreviewsByPromptId['p2']).toBeDefined() + expect(store.previewsByPromptId).toEqual({ p2: 'blob:b' }) + }) + + it('clears all previews', () => { + const store = useJobPreviewStore() + store.setPreviewUrl('p1', 'blob:a', 'node-1') + store.setPreviewUrl('p2', 'blob:b', 'node-2') + + store.clearAllPreviews() + + expect(store.nodePreviewsByPromptId).toEqual({}) + expect(store.previewsByPromptId).toEqual({}) + }) + + it('skips duplicate url', () => { + const store = useJobPreviewStore() + store.setPreviewUrl('p1', 'blob:a', 'node-1') + + store.setPreviewUrl('p1', 'blob:a', 'node-1') + + expect(releaseSharedObjectUrl).not.toHaveBeenCalled() + }) + + it('ignores setPreviewUrl when previews are disabled', () => { + previewMethodRef.value = 'none' + const store = useJobPreviewStore() + + store.setPreviewUrl('p1', 'blob:a', 'node-1') + + expect(store.nodePreviewsByPromptId).toEqual({}) + }) +}) diff --git a/src/stores/jobPreviewStore.ts b/src/stores/jobPreviewStore.ts index ce0b83a64a..f7dbe3c374 100644 --- a/src/stores/jobPreviewStore.ts +++ b/src/stores/jobPreviewStore.ts @@ -8,44 +8,59 @@ import { } from '@/utils/objectUrlUtil' type PromptPreviewMap = Record +interface NodePromptPreview { + url: string + nodeId?: string +} export const useJobPreviewStore = defineStore('jobPreview', () => { const settingStore = useSettingStore() - const previewsByPromptId = ref({}) - const readonlyPreviewsByPromptId = readonly(previewsByPromptId) + const nodePreviewsByPromptId = ref>({}) const previewMethod = computed(() => settingStore.get('Comfy.Execution.PreviewMethod') ) const isPreviewEnabled = computed(() => previewMethod.value !== 'none') - function setPreviewUrl(promptId: string | undefined, url: string) { + const previewsByPromptId = computed(() => { + const result: PromptPreviewMap = {} + for (const [k, v] of Object.entries(nodePreviewsByPromptId.value)) { + result[k] = v.url + } + return result + }) + + function setPreviewUrl( + promptId: string | undefined, + url: string, + nodeId?: string + ) { if (!promptId || !isPreviewEnabled.value) return - const current = previewsByPromptId.value[promptId] - if (current === url) return - if (current) releaseSharedObjectUrl(current) + const current = nodePreviewsByPromptId.value[promptId] + if (current?.url === url) return + if (current) releaseSharedObjectUrl(current.url) retainSharedObjectUrl(url) - previewsByPromptId.value = { - ...previewsByPromptId.value, - [promptId]: url + nodePreviewsByPromptId.value = { + ...nodePreviewsByPromptId.value, + [promptId]: { url, nodeId } } } function clearPreview(promptId: string | undefined) { if (!promptId) return - const current = previewsByPromptId.value[promptId] + const current = nodePreviewsByPromptId.value[promptId] if (!current) return - releaseSharedObjectUrl(current) - const next = { ...previewsByPromptId.value } + releaseSharedObjectUrl(current.url) + const next = { ...nodePreviewsByPromptId.value } delete next[promptId] - previewsByPromptId.value = next + nodePreviewsByPromptId.value = next } function clearAllPreviews() { - Object.values(previewsByPromptId.value).forEach((url) => { + for (const { url } of Object.values(nodePreviewsByPromptId.value)) { releaseSharedObjectUrl(url) - }) - previewsByPromptId.value = {} + } + nodePreviewsByPromptId.value = {} } watch(isPreviewEnabled, (enabled) => { @@ -53,7 +68,8 @@ export const useJobPreviewStore = defineStore('jobPreview', () => { }) return { - previewsByPromptId: readonlyPreviewsByPromptId, + nodePreviewsByPromptId: readonly(nodePreviewsByPromptId), + previewsByPromptId, isPreviewEnabled, setPreviewUrl, clearPreview,