From aabea4b78dcd258086f5c9c1fb008cf0683808f1 Mon Sep 17 00:00:00 2001 From: Christian Byrne Date: Wed, 30 Jul 2025 17:54:35 -0700 Subject: [PATCH] [feat] Viewport persistence for subgraph navigation (#4613) --- src/stores/subgraphNavigationStore.ts | 75 +++++- .../subgraphNavigationStore.viewport.test.ts | 239 ++++++++++++++++++ 2 files changed, 311 insertions(+), 3 deletions(-) create mode 100644 tests-ui/tests/store/subgraphNavigationStore.viewport.test.ts diff --git a/src/stores/subgraphNavigationStore.ts b/src/stores/subgraphNavigationStore.ts index f3e86e2ff..b52047128 100644 --- 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,6 +19,7 @@ export const useSubgraphNavigationStore = defineStore( 'subgraphNavigation', () => { const workflowStore = useWorkflowStore() + const canvasStore = useCanvasStore() /** The currently opened subgraph. */ const activeSubgraph = shallowRef() @@ -23,6 +27,11 @@ export const useSubgraphNavigationStore = defineStore( /** The stack of subgraph IDs from the root graph to the currently opened subgraph. */ const idStack = shallowReactive([]) + /** LRU cache for viewport states. Key: subgraph ID or 'root' for root graph */ + const viewportCache = new QuickLRU({ + 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,6 +137,9 @@ 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) } ) @@ -81,7 +147,10 @@ export const useSubgraphNavigationStore = defineStore( 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 new file mode 100644 index 000000000..d7439b679 --- /dev/null +++ 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) + }) + }) +})