fix: compositor cache survives workflow tab switches

Author: jtydhr88

src/extensions/core/imageCompositor.ts

@@ -1,11 +1,15 @@
 import type { CompositorBBox } from '@/renderer/extensions/compositor/composables/compositorLayerState'
 import { resetCompositorStateWidgets } from '@/renderer/extensions/compositor/composables/compositorWidgets'
+import type { CompositorLayersSnapshot } from '@/renderer/extensions/compositor/composables/useCompositorLayers'
 import {
   clearCompositorLayers,
   clearCompositorPreviewOverride,
-  setCompositorLayers
+  restoreCompositorLayers,
+  setCompositorLayers,
+  snapshotCompositorLayers
 } from '@/renderer/extensions/compositor/composables/useCompositorLayers'
 import type { LGraphNode } from '@/lib/litegraph/src/LGraphNode'
+import { registerWorkflowTransientState } from '@/platform/workflow/management/workflowTransientState'
 import type { NodeOutputWith } from '@/schemas/apiSchema'
 import { useExtensionService } from '@/services/extensionService'
 
@@ -16,6 +20,12 @@ type ImageCompositorOutput = NodeOutputWith<{
   compositor_state_stale?: boolean[]
 }>
 
+registerWorkflowTransientState('Comfy.ImageCompositor.layers', {
+  snapshot: snapshotCompositorLayers,
+  restore: (state) =>
+    restoreCompositorLayers(state as CompositorLayersSnapshot | undefined)
+})
+
 useExtensionService().registerExtension({
   name: 'Comfy.ImageCompositor',
 

src/platform/workflow/management/workflowTransientState.test.ts

@@ -0,0 +1,35 @@
+import { describe, expect, it, vi } from 'vitest'
+
+import {
+  registerWorkflowTransientState,
+  restoreWorkflowTransientState,
+  snapshotWorkflowTransientState
+} from './workflowTransientState'
+
+describe('workflowTransientState', () => {
+  it('round-trips each provider by its key', () => {
+    let stateA = 'a1'
+    registerWorkflowTransientState('test.a', {
+      snapshot: () => stateA,
+      restore: (state) => {
+        stateA = typeof state === 'string' ? state : 'cleared'
+      }
+    })
+
+    const snapshot = snapshotWorkflowTransientState()
+    stateA = 'a2'
+    restoreWorkflowTransientState(snapshot)
+    expect(stateA).toBe('a1')
+  })
+
+  it('restores undefined into providers missing from the snapshot', () => {
+    const restore = vi.fn()
+    registerWorkflowTransientState('test.b', { snapshot: () => 'b', restore })
+
+    restoreWorkflowTransientState(new Map())
+    expect(restore).toHaveBeenCalledWith(undefined)
+
+    restoreWorkflowTransientState(undefined)
+    expect(restore).toHaveBeenLastCalledWith(undefined)
+  })
+})

src/platform/workflow/management/workflowTransientState.ts

@@ -0,0 +1,29 @@
+export interface WorkflowTransientStateProvider {
+  snapshot(): unknown
+  restore(state: unknown): void
+}
+
+export type WorkflowTransientSnapshot = ReadonlyMap<string, unknown>
+
+const providers = new Map<string, WorkflowTransientStateProvider>()
+
+export function registerWorkflowTransientState(
+  key: string,
+  provider: WorkflowTransientStateProvider
+): void {
+  providers.set(key, provider)
+}
+
+export function snapshotWorkflowTransientState(): WorkflowTransientSnapshot {
+  return new Map(
+    [...providers].map(([key, provider]) => [key, provider.snapshot()])
+  )
+}
+
+export function restoreWorkflowTransientState(
+  snapshot: WorkflowTransientSnapshot | undefined
+): void {
+  for (const [key, provider] of providers) {
+    provider.restore(snapshot?.get(key))
+  }
+}

src/renderer/extensions/compositor/composables/useCompositorLayers.test.ts

@@ -11,8 +11,10 @@ import {
   getCompositorLayers,
   getCompositorPreviewOverride,
   hasCompositorLayers,
+  restoreCompositorLayers,
   setCompositorLayers,
-  setCompositorPreviewOverride
+  setCompositorPreviewOverride,
+  snapshotCompositorLayers
 } from './useCompositorLayers'
 
 const layerRef = { filename: 'a.png', subfolder: '', type: 'temp' }
@@ -97,6 +99,32 @@ describe('useCompositorLayers', () => {
     expect(getCompositorBBoxes(node)).toBeUndefined()
   })
 
+  it('survives a workflow switch through snapshot and restore', () => {
+    setCompositorLayers(node, [layerRef], ['hash-a'], [bbox])
+    const snapshot = snapshotCompositorLayers()
+
+    clearCompositorLayers(node)
+    expect(hasCompositorLayers(node)).toBe(false)
+
+    restoreCompositorLayers(snapshot)
+    expect(getCompositorLayers(node)).toEqual([layerRef])
+    expect(getCompositorInputsFingerprint(node)).toEqual(['hash-a'])
+    expect(getCompositorBBoxes(node)).toEqual([bbox])
+  })
+
+  it('restoring another workflow snapshot replaces the whole cache', () => {
+    setCompositorLayers(node, [layerRef])
+    const otherWorkflow = snapshotCompositorLayers()
+    setCompositorLayers(subgraphNode, [layerRef])
+
+    restoreCompositorLayers(otherWorkflow)
+    expect(hasCompositorLayers(node)).toBe(true)
+    expect(hasCompositorLayers(subgraphNode)).toBe(false)
+
+    restoreCompositorLayers(undefined)
+    expect(hasCompositorLayers(node)).toBe(false)
+  })
+
   it('stores, replaces and clears the preview override per node', () => {
     expect(getCompositorPreviewOverride(node)).toBeUndefined()
     setCompositorPreviewOverride(node, 'blob:first')

src/renderer/extensions/compositor/composables/useCompositorLayers.ts

@@ -15,6 +15,11 @@ interface CompositorNodeCache {
   bboxes?: (CompositorBBox | null)[]
 }
 
+export type CompositorLayersSnapshot = ReadonlyMap<
+  NodeLocatorId,
+  CompositorNodeCache
+>
+
 const cacheByNode = reactive(new Map<NodeLocatorId, CompositorNodeCache>())
 const previewOverrideByNode = reactive(new Map<NodeLocatorId, string>())
 
@@ -89,3 +94,16 @@ export function clearCompositorPreviewOverride(node: CompositorNodeRef): void {
 export function hasCompositorLayers(node: CompositorNodeRef): boolean {
   return (cacheByNode.get(cacheKey(node))?.layers.length ?? 0) > 0
 }
+
+export function snapshotCompositorLayers(): CompositorLayersSnapshot {
+  return new Map(cacheByNode)
+}
+
+export function restoreCompositorLayers(
+  snapshot: CompositorLayersSnapshot | undefined
+): void {
+  cacheByNode.clear()
+  for (const [key, value] of snapshot ?? []) {
+    cacheByNode.set(key, value)
+  }
+}

src/scripts/changeTracker.ts

@@ -3,6 +3,11 @@ import _ from 'es-toolkit/compat'
 
 import { assert } from '@/base/assert'
 import { LAYER_EDITOR_DIALOG_KEY } from '@/renderer/extensions/layerEditor/composables/layerEditorDialog'
+import type { WorkflowTransientSnapshot } from '@/platform/workflow/management/workflowTransientState'
+import {
+  restoreWorkflowTransientState,
+  snapshotWorkflowTransientState
+} from '@/platform/workflow/management/workflowTransientState'
 import type { CanvasPointerEvent } from '@/lib/litegraph/src/litegraph'
 import { LGraphCanvas, LiteGraph } from '@/lib/litegraph/src/litegraph'
 import type { ComfyWorkflow } from '@/platform/workflow/management/stores/workflowStore'
@@ -269,6 +274,7 @@ export class ChangeTracker {
 
   ds?: { scale: number; offset: [number, number] }
   nodeOutputs?: Record<string, ExecutedWsMessage['output']>
+  private transientState?: WorkflowTransientSnapshot
 
   private subgraphState?: {
     navigation: string[]
@@ -304,6 +310,7 @@ export class ChangeTracker {
       offset: [app.canvas.ds.offset[0], app.canvas.ds.offset[1]]
     }
     this.nodeOutputs = useNodeOutputStore().snapshotOutputs()
+    this.transientState = snapshotWorkflowTransientState()
     const navigation = useSubgraphNavigationStore().exportState()
     // Always store the navigation state, even if empty (root level)
     this.subgraphState = { navigation }
@@ -349,6 +356,7 @@ export class ChangeTracker {
     if (this.nodeOutputs) {
       useNodeOutputStore().restoreOutputs(this.nodeOutputs)
     }
+    restoreWorkflowTransientState(this.transientState)
     if (this.subgraphState) {
       const { navigation } = this.subgraphState
       const firstInvalidIndex = navigation.findIndex(