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

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

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