[feat] Viewport persistence for subgraph navigation (#4613)

christian-byrne · web-flow · commit aabea4b78dcd · 2025-07-30T17:54:35.000-07:00
diff --git a/src/stores/subgraphNavigationStore.ts b/src/stores/subgraphNavigationStore.ts
@@ -1,10 +1,13 @@
+import QuickLRU from '@alloc/quick-lru'
 import type { Subgraph } from '@comfyorg/litegraph'
+import type { DragAndScaleState } from '@comfyorg/litegraph/dist/DragAndScale'
 import { defineStore } from 'pinia'
 import { computed, shallowReactive, shallowRef, watch } from 'vue'
 
 import { app } from '@/scripts/app'
 import { isNonNullish } from '@/utils/typeGuardUtil'
 
+import { useCanvasStore } from './graphStore'
 import { useWorkflowStore } from './workflowStore'
 
 /**
@@ -16,13 +19,19 @@ export const useSubgraphNavigationStore = defineStore(
   'subgraphNavigation',
   () => {
     const workflowStore = useWorkflowStore()
+    const canvasStore = useCanvasStore()
 
     /** The currently opened subgraph. */
     const activeSubgraph = shallowRef<Subgraph>()
 
     /** The stack of subgraph IDs from the root graph to the currently opened subgraph. */
     const idStack = shallowReactive<string[]>([])
 
+    /** LRU cache for viewport states. Key: subgraph ID or 'root' for root graph */
+    const viewportCache = new QuickLRU<string, DragAndScaleState>({
+      maxSize: 32
+    })
+
     /**
      * A stack representing subgraph navigation history from the root graph to
      * the current opened subgraph.
@@ -48,19 +57,73 @@ export const useSubgraphNavigationStore = defineStore(
      */
     const exportState = () => [...idStack]
 
+    /**
+     * Get the current viewport state.
+     * @returns The current viewport state, or null if the canvas is not available.
+     */
+    const getCurrentViewport = (): DragAndScaleState | null => {
+      const canvas = canvasStore.getCanvas()
+      if (!canvas) return null
+
+      return {
+        scale: canvas.ds.state.scale,
+        offset: [...canvas.ds.state.offset]
+      }
+    }
+
+    /**
+     * Save the current viewport state.
+     * @param graphId The graph ID to save for. Use 'root' for root graph, or omit to use current context.
+     */
+    const saveViewport = (graphId: string) => {
+      const viewport = getCurrentViewport()
+      if (!viewport) return
+
+      viewportCache.set(graphId, viewport)
+    }
+
+    /**
+     * Restore viewport state for a graph.
+     * @param graphId The graph ID to restore. Use 'root' for root graph, or omit to use current context.
+     */
+    const restoreViewport = (graphId: string) => {
+      const viewport = viewportCache.get(graphId)
+      if (!viewport) return
+
+      const canvas = app.canvas
+      if (!canvas) return
+
+      canvas.ds.scale = viewport.scale
+      canvas.ds.offset[0] = viewport.offset[0]
+      canvas.ds.offset[1] = viewport.offset[1]
+      canvas.setDirty(true, true)
+    }
+
     // Reset on workflow change
     watch(
       () => workflowStore.activeWorkflow,
-      () => (idStack.length = 0)
+      () => {
+        idStack.length = 0
+      }
     )
 
     // Update navigation stack when opened subgraph changes
     watch(
       () => workflowStore.activeSubgraph,
-      (subgraph) => {
+      (subgraph, prevSubgraph) => {
+        // Save viewport state for the graph we're leaving
+        if (prevSubgraph) {
+          // Leaving a subgraph
+          saveViewport(prevSubgraph.id)
+        } else if (!prevSubgraph && subgraph) {
+          // Leaving root graph to enter a subgraph
+          saveViewport('root')
+        }
+
         // Navigated back to the root graph
         if (!subgraph) {
           idStack.length = 0
+          restoreViewport('root')
           return
         }
 
@@ -74,14 +137,20 @@ export const useSubgraphNavigationStore = defineStore(
           // Navigated to a different subgraph
           idStack.splice(index + 1, lastIndex - index)
         }
+
+        // Always try to restore viewport for the target subgraph
+        restoreViewport(subgraph.id)
       }
     )
 
     return {
       activeSubgraph,
       navigationStack,
       restoreState,
-      exportState
+      exportState,
+      saveViewport,
+      restoreViewport,
+      viewportCache
     }
   }
 )
diff --git a/tests-ui/tests/store/subgraphNavigationStore.viewport.test.ts b/tests-ui/tests/store/subgraphNavigationStore.viewport.test.ts
@@ -0,0 +1,239 @@
+import { createPinia, setActivePinia } from 'pinia'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { nextTick } from 'vue'
+
+import { app } from '@/scripts/app'
+import { useSubgraphNavigationStore } from '@/stores/subgraphNavigationStore'
+import { useWorkflowStore } from '@/stores/workflowStore'
+import type { ComfyWorkflow } from '@/stores/workflowStore'
+
+vi.mock('@/scripts/app', () => {
+  const mockCanvas = {
+    subgraph: null,
+    ds: {
+      scale: 1,
+      offset: [0, 0],
+      state: {
+        scale: 1,
+        offset: [0, 0]
+      }
+    },
+    setDirty: vi.fn()
+  }
+
+  return {
+    app: {
+      graph: {
+        subgraphs: new Map(),
+        getNodeById: vi.fn()
+      },
+      canvas: mockCanvas
+    }
+  }
+})
+
+// Mock canvasStore
+vi.mock('@/stores/graphStore', () => ({
+  useCanvasStore: () => ({
+    getCanvas: () => (app as any).canvas
+  })
+}))
+
+// Get reference to mock canvas
+const mockCanvas = app.canvas as any
+
+describe('useSubgraphNavigationStore - Viewport Persistence', () => {
+  beforeEach(() => {
+    setActivePinia(createPinia())
+    // Reset canvas state
+    mockCanvas.ds.scale = 1
+    mockCanvas.ds.offset = [0, 0]
+    mockCanvas.ds.state.scale = 1
+    mockCanvas.ds.state.offset = [0, 0]
+    mockCanvas.setDirty.mockClear()
+  })
+
+  describe('saveViewport', () => {
+    it('should save viewport state for root graph', () => {
+      const navigationStore = useSubgraphNavigationStore()
+
+      // Set viewport state
+      mockCanvas.ds.state.scale = 2
+      mockCanvas.ds.state.offset = [100, 200]
+
+      // Save viewport for root
+      navigationStore.saveViewport('root')
+
+      // Check it was saved
+      const saved = navigationStore.viewportCache.get('root')
+      expect(saved).toEqual({
+        scale: 2,
+        offset: [100, 200]
+      })
+    })
+
+    it('should save viewport state for subgraph', () => {
+      const navigationStore = useSubgraphNavigationStore()
+
+      // Set viewport state
+      mockCanvas.ds.state.scale = 1.5
+      mockCanvas.ds.state.offset = [50, 75]
+
+      // Save viewport for subgraph
+      navigationStore.saveViewport('subgraph-123')
+
+      // Check it was saved
+      const saved = navigationStore.viewportCache.get('subgraph-123')
+      expect(saved).toEqual({
+        scale: 1.5,
+        offset: [50, 75]
+      })
+    })
+
+    it('should save viewport for current context when no ID provided', () => {
+      const navigationStore = useSubgraphNavigationStore()
+      const workflowStore = useWorkflowStore()
+
+      // Mock being in a subgraph
+      const mockSubgraph = { id: 'sub-456' }
+      workflowStore.activeSubgraph = mockSubgraph as any
+
+      // Set viewport state
+      mockCanvas.ds.state.scale = 3
+      mockCanvas.ds.state.offset = [10, 20]
+
+      // Save viewport without ID (should default to root since activeSubgraph is not tracked by navigation store)
+      navigationStore.saveViewport('sub-456')
+
+      // Should save for the specified subgraph
+      const saved = navigationStore.viewportCache.get('sub-456')
+      expect(saved).toEqual({
+        scale: 3,
+        offset: [10, 20]
+      })
+    })
+  })
+
+  describe('restoreViewport', () => {
+    it('should restore viewport state for root graph', () => {
+      const navigationStore = useSubgraphNavigationStore()
+
+      // Save a viewport state
+      navigationStore.viewportCache.set('root', {
+        scale: 2.5,
+        offset: [150, 250]
+      })
+
+      // Restore it
+      navigationStore.restoreViewport('root')
+
+      // Check canvas was updated
+      expect(mockCanvas.ds.scale).toBe(2.5)
+      expect(mockCanvas.ds.offset).toEqual([150, 250])
+      expect(mockCanvas.setDirty).toHaveBeenCalledWith(true, true)
+    })
+
+    it('should restore viewport state for subgraph', () => {
+      const navigationStore = useSubgraphNavigationStore()
+
+      // Save a viewport state
+      navigationStore.viewportCache.set('sub-789', {
+        scale: 0.75,
+        offset: [-50, -100]
+      })
+
+      // Restore it
+      navigationStore.restoreViewport('sub-789')
+
+      // Check canvas was updated
+      expect(mockCanvas.ds.scale).toBe(0.75)
+      expect(mockCanvas.ds.offset).toEqual([-50, -100])
+    })
+
+    it('should do nothing if no saved viewport exists', () => {
+      const navigationStore = useSubgraphNavigationStore()
+
+      // Reset canvas
+      mockCanvas.ds.scale = 1
+      mockCanvas.ds.offset = [0, 0]
+      mockCanvas.setDirty.mockClear()
+
+      // Try to restore non-existent viewport
+      navigationStore.restoreViewport('non-existent')
+
+      // Canvas should not change
+      expect(mockCanvas.ds.scale).toBe(1)
+      expect(mockCanvas.ds.offset).toEqual([0, 0])
+      expect(mockCanvas.setDirty).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('navigation integration', () => {
+    it('should save and restore viewport when navigating between subgraphs', async () => {
+      const navigationStore = useSubgraphNavigationStore()
+      const workflowStore = useWorkflowStore()
+
+      // Start at root with custom viewport
+      mockCanvas.ds.state.scale = 2
+      mockCanvas.ds.state.offset = [100, 100]
+
+      // Navigate to subgraph
+      const subgraph1 = { id: 'sub1' }
+      workflowStore.activeSubgraph = subgraph1 as any
+      await nextTick()
+
+      // Root viewport should have been saved
+      const rootViewport = navigationStore.viewportCache.get('root')
+      expect(rootViewport).toBeDefined()
+      expect(rootViewport?.scale).toBe(2)
+      expect(rootViewport?.offset).toEqual([100, 100])
+
+      // Change viewport in subgraph
+      mockCanvas.ds.state.scale = 0.5
+      mockCanvas.ds.state.offset = [-50, -50]
+
+      // Navigate back to root
+      workflowStore.activeSubgraph = undefined
+      await nextTick()
+
+      // Subgraph viewport should have been saved
+      const sub1Viewport = navigationStore.viewportCache.get('sub1')
+      expect(sub1Viewport).toBeDefined()
+      expect(sub1Viewport?.scale).toBe(0.5)
+      expect(sub1Viewport?.offset).toEqual([-50, -50])
+
+      // Root viewport should be restored
+      expect(mockCanvas.ds.scale).toBe(2)
+      expect(mockCanvas.ds.offset).toEqual([100, 100])
+    })
+
+    it('should preserve viewport cache when switching workflows', async () => {
+      const navigationStore = useSubgraphNavigationStore()
+      const workflowStore = useWorkflowStore()
+
+      // Add some viewport states
+      navigationStore.viewportCache.set('root', { scale: 2, offset: [0, 0] })
+      navigationStore.viewportCache.set('sub1', {
+        scale: 1.5,
+        offset: [10, 10]
+      })
+
+      expect(navigationStore.viewportCache.size).toBe(2)
+
+      // Switch workflows
+      const workflow1 = { path: 'workflow1.json' } as ComfyWorkflow
+      const workflow2 = { path: 'workflow2.json' } as ComfyWorkflow
+
+      workflowStore.activeWorkflow = workflow1 as any
+      await nextTick()
+
+      workflowStore.activeWorkflow = workflow2 as any
+      await nextTick()
+
+      // Cache should be preserved (LRU will manage memory)
+      expect(navigationStore.viewportCache.size).toBe(2)
+      expect(navigationStore.viewportCache.has('root')).toBe(true)
+      expect(navigationStore.viewportCache.has('sub1')).toBe(true)
+    })
+  })
+})
