From 46d964f520214d120d64182015952d45af68b3be Mon Sep 17 00:00:00 2001
From: filtered <176114999+webfiltered@users.noreply.github.com>
Date: Thu, 15 May 2025 14:23:57 +1000
Subject: [PATCH] Keep subgraph nav state when swapping workflows

---
 src/scripts/changeTracker.ts          | 17 +++++++++++
 src/stores/subgraphNavigationStore.ts | 44 ++++++++++++++++++---------
 2 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/src/scripts/changeTracker.ts b/src/scripts/changeTracker.ts
index 2186043ed..8eeed1de4 100644
--- a/src/scripts/changeTracker.ts
+++ b/src/scripts/changeTracker.ts
@@ -6,6 +6,7 @@ import log from 'loglevel'
 import type { ExecutedWsMessage } from '@/schemas/apiSchema'
 import type { ComfyWorkflowJSON } from '@/schemas/comfyWorkflowSchema'
 import { useExecutionStore } from '@/stores/executionStore'
+import { useSubgraphNavigationStore } from '@/stores/subgraphNavigationStore'
 import { ComfyWorkflow, useWorkflowStore } from '@/stores/workflowStore'
 
 import { api } from './api'
@@ -37,6 +38,10 @@ export class ChangeTracker {
   ds?: { scale: number; offset: [number, number] }
   nodeOutputs?: Record<string, any>
 
+  private subgraphState?: {
+    navigation: string[]
+  }
+
   constructor(
     /**
      * The workflow that this change tracker is tracking
@@ -67,6 +72,8 @@ export class ChangeTracker {
       scale: app.canvas.ds.scale,
       offset: [app.canvas.ds.offset[0], app.canvas.ds.offset[1]]
     }
+    const navigation = useSubgraphNavigationStore().exportState()
+    this.subgraphState = navigation.length ? { navigation } : undefined
   }
 
   restore() {
@@ -77,6 +84,16 @@ export class ChangeTracker {
     if (this.nodeOutputs) {
       app.nodeOutputs = this.nodeOutputs
     }
+    if (this.subgraphState) {
+      const { navigation } = this.subgraphState
+      useSubgraphNavigationStore().restoreState(navigation)
+
+      const activeId = navigation.at(-1)
+      if (activeId) {
+        const subgraph = app.graph.subgraphs.get(activeId)
+        if (subgraph) app.canvas.setGraph(subgraph)
+      }
+    }
   }
 
   updateModified() {
diff --git a/src/stores/subgraphNavigationStore.ts b/src/stores/subgraphNavigationStore.ts
index d07676337..a09ad400e 100644
--- a/src/stores/subgraphNavigationStore.ts
+++ b/src/stores/subgraphNavigationStore.ts
@@ -1,10 +1,10 @@
 import type { Subgraph } from '@comfyorg/litegraph'
 import { defineStore } from 'pinia'
-import { computed, shallowReactive, shallowRef, watch } from 'vue'
+import { computed, shallowRef, watch } from 'vue'
 
+import { app } from '@/scripts/app'
 import { isNonNullish } from '@/utils/typeGuardUtil'
 
-import { useCanvasStore } from './graphStore'
 import { useWorkflowStore } from './workflowStore'
 
 /**
@@ -16,28 +16,42 @@ 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 subgraphIdStack = shallowReactive<string[]>([])
+    const idStack: string[] = []
 
     /**
      * A stack representing subgraph navigation history from the root graph to
      * the current opened subgraph.
      */
     const navigationStack = computed(() =>
-      subgraphIdStack
-        .map((id) => canvasStore.getCanvas().graph?.subgraphs.get(id))
-        .filter(isNonNullish)
+      idStack.map((id) => app.graph.subgraphs.get(id)).filter(isNonNullish)
     )
 
+    /**
+     * Restore the navigation stack from a list of subgraph IDs.
+     * @param subgraphIds The list of subgraph IDs to restore the navigation stack from.
+     * @see exportState
+     */
+    const restoreState = (subgraphIds: string[]) => {
+      idStack.length = 0
+      for (const id of subgraphIds) idStack.push(id)
+    }
+
+    /**
+     * Export the navigation stack as a list of subgraph IDs.
+     * @returns The list of subgraph IDs, ending with the currently active subgraph.
+     * @see restoreState
+     */
+    const exportState = () => [...idStack]
+
     // Reset on workflow change
     watch(
       () => workflowStore.activeWorkflow,
-      () => (subgraphIdStack.length = 0)
+      () => (idStack.length = 0)
     )
 
     // Update navigation stack when opened subgraph changes
@@ -46,26 +60,28 @@ export const useSubgraphNavigationStore = defineStore(
       (subgraph) => {
         // Navigated back to the root graph
         if (!subgraph) {
-          subgraphIdStack.length = 0
+          idStack.length = 0
           return
         }
 
-        const index = subgraphIdStack.lastIndexOf(subgraph.id)
-        const lastIndex = subgraphIdStack.length - 1
+        const index = idStack.lastIndexOf(subgraph.id)
+        const lastIndex = idStack.length - 1
 
         if (index === -1) {
           // Opened a new subgraph
-          subgraphIdStack.push(subgraph.id)
+          idStack.push(subgraph.id)
         } else if (index !== lastIndex) {
           // Navigated to a different subgraph
-          subgraphIdStack.splice(index + 1, lastIndex - index)
+          idStack.splice(index + 1, lastIndex - index)
         }
       }
     )
 
     return {
       activeSubgraph,
-      navigationStack
+      navigationStack,
+      restoreState,
+      exportState
     }
   }
 )