diff --git a/src/composables/useMinimap.ts b/src/composables/useMinimap.ts index 7ec25c585..f1b85013b 100644 --- a/src/composables/useMinimap.ts +++ b/src/composables/useMinimap.ts @@ -115,6 +115,21 @@ export function useMinimap() { () => workflowStore.activeSubgraph, () => { graph.value = app.canvas?.graph + // Force viewport update when switching subgraphs + if (initialized.value && visible.value) { + updateViewport() + } + } + ) + + // Update viewport when switching workflows + watch( + () => workflowStore.activeWorkflow, + () => { + // Force viewport update when switching workflows + if (initialized.value && visible.value) { + updateViewport() + } } ) diff --git a/src/scripts/changeTracker.ts b/src/scripts/changeTracker.ts index f6a03afd5..ee7ccae0c 100644 --- a/src/scripts/changeTracker.ts +++ b/src/scripts/changeTracker.ts @@ -73,7 +73,8 @@ export class ChangeTracker { offset: [app.canvas.ds.offset[0], app.canvas.ds.offset[1]] } const navigation = useSubgraphNavigationStore().exportState() - this.subgraphState = navigation.length ? { navigation } : undefined + // Always store the navigation state, even if empty (root level) + this.subgraphState = { navigation } } restore() { @@ -90,8 +91,14 @@ export class ChangeTracker { const activeId = navigation.at(-1) if (activeId) { + // Navigate to the saved subgraph const subgraph = app.graph.subgraphs.get(activeId) - if (subgraph) app.canvas.setGraph(subgraph) + if (subgraph) { + app.canvas.setGraph(subgraph) + } + } else { + // Empty navigation array means root level + app.canvas.setGraph(app.graph) } } } diff --git a/src/stores/subgraphNavigationStore.ts b/src/stores/subgraphNavigationStore.ts index b52047128..677c3d2fb 100644 --- a/src/stores/subgraphNavigationStore.ts +++ b/src/stores/subgraphNavigationStore.ts @@ -2,9 +2,10 @@ 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 { computed, ref, shallowRef, watch } from 'vue' import { app } from '@/scripts/app' +import { findSubgraphPathById } from '@/utils/graphTraversalUtil' import { isNonNullish } from '@/utils/typeGuardUtil' import { useCanvasStore } from './graphStore' @@ -25,7 +26,7 @@ export const useSubgraphNavigationStore = defineStore( const activeSubgraph = shallowRef() /** The stack of subgraph IDs from the root graph to the currently opened subgraph. */ - const idStack = shallowReactive([]) + const idStack = ref([]) /** LRU cache for viewport states. Key: subgraph ID or 'root' for root graph */ const viewportCache = new QuickLRU({ @@ -37,7 +38,9 @@ export const useSubgraphNavigationStore = defineStore( * the current opened subgraph. */ const navigationStack = computed(() => - idStack.map((id) => app.graph.subgraphs.get(id)).filter(isNonNullish) + idStack.value + .map((id) => app.graph.subgraphs.get(id)) + .filter(isNonNullish) ) /** @@ -46,8 +49,8 @@ export const useSubgraphNavigationStore = defineStore( * @see exportState */ const restoreState = (subgraphIds: string[]) => { - idStack.length = 0 - for (const id of subgraphIds) idStack.push(id) + idStack.value.length = 0 + for (const id of subgraphIds) idStack.value.push(id) } /** @@ -55,7 +58,7 @@ export const useSubgraphNavigationStore = defineStore( * @returns The list of subgraph IDs, ending with the currently active subgraph. * @see restoreState */ - const exportState = () => [...idStack] + const exportState = () => [...idStack.value] /** * Get the current viewport state. @@ -99,47 +102,49 @@ export const useSubgraphNavigationStore = defineStore( canvas.setDirty(true, true) } - // Reset on workflow change - watch( - () => workflowStore.activeWorkflow, - () => { - idStack.length = 0 + /** + * Update the navigation stack when the active subgraph changes. + * @param subgraph The new active subgraph. + * @param prevSubgraph The previous active subgraph. + */ + const onNavigated = ( + subgraph: Subgraph | undefined, + prevSubgraph: Subgraph | undefined + ) => { + // 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') } - ) - // Update navigation stack when opened subgraph changes + const isInRootGraph = !subgraph + if (isInRootGraph) { + idStack.value.length = 0 + restoreViewport('root') + return + } + + const path = findSubgraphPathById(subgraph.rootGraph, subgraph.id) + const isInReachableSubgraph = !!path + if (isInReachableSubgraph) { + idStack.value = [...path] + } else { + // Treat as if opening a new subgraph + idStack.value = [subgraph.id] + } + + // Always try to restore viewport for the target subgraph + restoreViewport(subgraph.id) + } + + // Update navigation stack when opened subgraph changes (also triggers when switching workflows) watch( () => workflowStore.activeSubgraph, - (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 - } - - const index = idStack.lastIndexOf(subgraph.id) - const lastIndex = idStack.length - 1 - - if (index === -1) { - // Opened a new subgraph - idStack.push(subgraph.id) - } else if (index !== lastIndex) { - // Navigated to a different subgraph - idStack.splice(index + 1, lastIndex - index) - } - - // Always try to restore viewport for the target subgraph - restoreViewport(subgraph.id) + (newValue, oldValue) => { + onNavigated(newValue, oldValue) } ) diff --git a/src/utils/graphTraversalUtil.ts b/src/utils/graphTraversalUtil.ts index 7124c4a60..3adb87726 100644 --- a/src/utils/graphTraversalUtil.ts +++ b/src/utils/graphTraversalUtil.ts @@ -205,7 +205,7 @@ export function findSubgraphByUuid( targetUuid: string ): Subgraph | null { // Check all nodes in the current graph - for (const node of graph._nodes) { + for (const node of graph.nodes) { if (node.isSubgraphNode?.() && node.subgraph) { if (node.subgraph.id === targetUuid) { return node.subgraph @@ -218,6 +218,42 @@ export function findSubgraphByUuid( return null } +/** + * Iteratively finds the path of subgraph IDs to a target subgraph. + * @param rootGraph The graph to start searching from. + * @param targetId The ID of the subgraph to find. + * @returns An array of subgraph IDs representing the path, or `null` if not found. + */ +export function findSubgraphPathById( + rootGraph: LGraph, + targetId: string +): string[] | null { + const stack: { graph: LGraph | Subgraph; path: string[] }[] = [ + { graph: rootGraph, path: [] } + ] + + while (stack.length > 0) { + const { graph, path } = stack.pop()! + + // Check if graph exists and has _nodes property + if (!graph || !graph._nodes || !Array.isArray(graph._nodes)) { + continue + } + + for (const node of graph._nodes) { + if (node.isSubgraphNode?.() && node.subgraph) { + const newPath = [...path, String(node.subgraph.id)] + if (node.subgraph.id === targetId) { + return newPath + } + stack.push({ graph: node.subgraph, path: newPath }) + } + } + } + + return null +} + /** * Get a node by its execution ID from anywhere in the graph hierarchy. * Execution IDs use hierarchical format like "123:456:789" for nested nodes. diff --git a/tests-ui/tests/store/subgraphNavigationStore.test.ts b/tests-ui/tests/store/subgraphNavigationStore.test.ts index 4ee63353b..b9d5537cd 100644 --- a/tests-ui/tests/store/subgraphNavigationStore.test.ts +++ b/tests-ui/tests/store/subgraphNavigationStore.test.ts @@ -2,20 +2,46 @@ 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', () => ({ - app: { - graph: { - subgraphs: new Map(), - getNodeById: vi.fn() +vi.mock('@/scripts/app', () => { + const mockCanvas = { + subgraph: null, + ds: { + scale: 1, + offset: [0, 0], + state: { + scale: 1, + offset: [0, 0] + } }, - canvas: { - subgraph: null + setDirty: vi.fn() + } + + return { + app: { + graph: { + _nodes: [], + nodes: [], + subgraphs: new Map(), + getNodeById: vi.fn() + }, + canvas: mockCanvas } } +}) + +vi.mock('@/stores/graphStore', () => ({ + useCanvasStore: () => ({ + getCanvas: () => (app as any).canvas + }) +})) + +vi.mock('@/utils/graphTraversalUtil', () => ({ + findSubgraphPathById: vi.fn() })) describe('useSubgraphNavigationStore', () => { @@ -51,42 +77,77 @@ describe('useSubgraphNavigationStore', () => { expect(navigationStore.exportState()).toEqual(['subgraph-1', 'subgraph-2']) }) - it('should clear navigation stack when switching to a different workflow', async () => { + it('should preserve navigation stack per workflow', async () => { const navigationStore = useSubgraphNavigationStore() const workflowStore = useWorkflowStore() // Mock first workflow const workflow1 = { path: 'workflow1.json', - filename: 'workflow1.json' - } as ComfyWorkflow + filename: 'workflow1.json', + changeTracker: { + restore: vi.fn(), + store: vi.fn() + } + } as unknown as ComfyWorkflow // Set the active workflow workflowStore.activeWorkflow = workflow1 as any - // Simulate being in a subgraph + // Simulate the restore process that happens when loading a workflow + // Since subgraphState is private, we'll simulate the effect by directly restoring navigation navigationStore.restoreState(['subgraph-1', 'subgraph-2']) + // Verify navigation was set expect(navigationStore.exportState()).toHaveLength(2) + expect(navigationStore.exportState()).toEqual(['subgraph-1', 'subgraph-2']) - // Switch to a different workflow + // Switch to a different workflow with no subgraph state (root level) const workflow2 = { path: 'workflow2.json', - filename: 'workflow2.json' - } as ComfyWorkflow + filename: 'workflow2.json', + changeTracker: { + restore: vi.fn(), + store: vi.fn() + } + } as unknown as ComfyWorkflow workflowStore.activeWorkflow = workflow2 as any - // Wait for Vue's reactivity to process the change - await nextTick() + // Simulate the restore process for workflow2 + // Since subgraphState is private, we'll simulate the effect by directly restoring navigation + navigationStore.restoreState([]) - // The navigation stack SHOULD be cleared because we switched workflows + // The navigation stack should be empty for workflow2 (at root level) expect(navigationStore.exportState()).toHaveLength(0) + + // Switch back to workflow1 + workflowStore.activeWorkflow = workflow1 as any + + // Simulate the restore process for workflow1 again + // Since subgraphState is private, we'll simulate the effect by directly restoring navigation + navigationStore.restoreState(['subgraph-1', 'subgraph-2']) + + // The navigation stack should be restored for workflow1 + expect(navigationStore.exportState()).toHaveLength(2) + expect(navigationStore.exportState()).toEqual(['subgraph-1', 'subgraph-2']) }) - it('should handle null workflow gracefully', async () => { + it('should clear navigation when activeSubgraph becomes undefined', async () => { const navigationStore = useSubgraphNavigationStore() const workflowStore = useWorkflowStore() + const { findSubgraphPathById } = await import('@/utils/graphTraversalUtil') + + // Create mock subgraph and graph structure + const mockSubgraph = { + id: 'subgraph-1', + rootGraph: (app as any).graph, + _nodes: [], + nodes: [] + } + + // Add the subgraph to the graph's subgraphs map + ;(app as any).graph.subgraphs.set('subgraph-1', mockSubgraph) // First set an active workflow const mockWorkflow = { @@ -95,19 +156,29 @@ describe('useSubgraphNavigationStore', () => { } as ComfyWorkflow workflowStore.activeWorkflow = mockWorkflow as any - await nextTick() - // Add some items to the navigation stack - navigationStore.restoreState(['subgraph-1']) - expect(navigationStore.exportState()).toHaveLength(1) + // Mock findSubgraphPathById to return the correct path + vi.mocked(findSubgraphPathById).mockReturnValue(['subgraph-1']) - // Set workflow to null - workflowStore.activeWorkflow = null + // Set canvas.subgraph and trigger update to set activeSubgraph + ;(app as any).canvas.subgraph = mockSubgraph + workflowStore.updateActiveGraph() // Wait for Vue's reactivity to process the change await nextTick() - // Stack should be cleared when workflow becomes null + // Verify navigation was set by the watcher + expect(navigationStore.exportState()).toHaveLength(1) + expect(navigationStore.exportState()).toEqual(['subgraph-1']) + + // Clear canvas.subgraph and trigger update (simulating navigating back to root) + ;(app as any).canvas.subgraph = null + workflowStore.updateActiveGraph() + + // Wait for Vue's reactivity to process the change + await nextTick() + + // Stack should be cleared when activeSubgraph becomes undefined expect(navigationStore.exportState()).toHaveLength(0) }) }) diff --git a/tests-ui/tests/store/subgraphNavigationStore.viewport.test.ts b/tests-ui/tests/store/subgraphNavigationStore.viewport.test.ts index d7439b679..db40609e0 100644 --- a/tests-ui/tests/store/subgraphNavigationStore.viewport.test.ts +++ b/tests-ui/tests/store/subgraphNavigationStore.viewport.test.ts @@ -24,6 +24,8 @@ vi.mock('@/scripts/app', () => { return { app: { graph: { + _nodes: [], + nodes: [], subgraphs: new Map(), getNodeById: vi.fn() }, @@ -173,16 +175,29 @@ describe('useSubgraphNavigationStore - Viewport Persistence', () => { const navigationStore = useSubgraphNavigationStore() const workflowStore = useWorkflowStore() + // Create mock subgraph with both _nodes and nodes properties + const mockRootGraph = { + _nodes: [], + nodes: [], + subgraphs: new Map(), + getNodeById: vi.fn() + } + const subgraph1 = { + id: 'sub1', + rootGraph: mockRootGraph, + _nodes: [], + nodes: [] + } + // 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 + // Root viewport should have been saved automatically const rootViewport = navigationStore.viewportCache.get('root') expect(rootViewport).toBeDefined() expect(rootViewport?.scale).toBe(2) @@ -196,13 +211,13 @@ describe('useSubgraphNavigationStore - Viewport Persistence', () => { workflowStore.activeSubgraph = undefined await nextTick() - // Subgraph viewport should have been saved + // Subgraph viewport should have been saved automatically 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 + // Root viewport should be restored automatically expect(mockCanvas.ds.scale).toBe(2) expect(mockCanvas.ds.offset).toEqual([100, 100]) }) diff --git a/tests-ui/tests/store/workflowStore.test.ts b/tests-ui/tests/store/workflowStore.test.ts index 579f330ff..a82e1adfc 100644 --- a/tests-ui/tests/store/workflowStore.test.ts +++ b/tests-ui/tests/store/workflowStore.test.ts @@ -597,7 +597,9 @@ describe('useWorkflowStore', () => { // Setup mock graph structure with subgraphs const mockSubgraph = { id: 'a1b2c3d4-e5f6-7890-abcd-ef1234567890', - _nodes: [] + rootGraph: null as any, + _nodes: [], + nodes: [] } const mockNode = { @@ -608,6 +610,7 @@ describe('useWorkflowStore', () => { const mockRootGraph = { _nodes: [mockNode], + nodes: [mockNode], subgraphs: new Map([[mockSubgraph.id, mockSubgraph]]), getNodeById: (id: string | number) => { if (String(id) === '123') return mockNode @@ -615,6 +618,8 @@ describe('useWorkflowStore', () => { } } + mockSubgraph.rootGraph = mockRootGraph as any + vi.mocked(comfyApp).graph = mockRootGraph as any vi.mocked(comfyApp.canvas).subgraph = mockSubgraph as any store.activeSubgraph = mockSubgraph as any