[fix] Preserve per-workflow subgraph navigation state (#4616)

Co-authored-by: Claude <noreply@anthropic.com>

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()
+      }
     }
   )
 

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)
       }
     }
   }

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<Subgraph>()
 
     /** The stack of subgraph IDs from the root graph to the currently opened subgraph. */
-    const idStack = shallowReactive<string[]>([])
+    const idStack = ref<string[]>([])
 
     /** LRU cache for viewport states. Key: subgraph ID or 'root' for root graph */
     const viewportCache = new QuickLRU<string, DragAndScaleState>({
@@ -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)
       }
     )
 

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.

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)
   })
 })

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])
     })

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