ComfyUI_frontend/src/stores/subgraphNavigationStore.ts

import QuickLRU from '@alloc/quick-lru'
import { useRouteHash } from '@vueuse/router'
import { defineStore } from 'pinia'
import { computed, ref, shallowRef, watch } from 'vue'
import {
  NavigationFailureType,
  isNavigationFailure,
  useRouter
} from 'vue-router'

import type { DragAndScaleState } from '@/lib/litegraph/src/DragAndScale'
import type { Subgraph } from '@/lib/litegraph/src/litegraph'
import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
import { useWorkflowService } from '@/platform/workflow/core/services/workflowService'
import { useCanvasStore } from '@/renderer/core/canvas/canvasStore'
import { requestSlotLayoutSyncForAllNodes } from '@/renderer/extensions/vueNodes/composables/useSlotElementTracking'
import { isUuidShapedSubgraphId } from '@/schemas/subgraphIdSchema'
import { app } from '@/scripts/app'
import { useLitegraphService } from '@/services/litegraphService'
import { findSubgraphPathById } from '@/utils/graphTraversalUtil'
import { isNonNullish, isSubgraph } from '@/utils/typeGuardUtil'

export const VIEWPORT_CACHE_MAX_SIZE = 32

/**
 * Stores the current subgraph navigation state; a stack representing subgraph
 * navigation history from the root graph to the subgraph that is currently
 * open.
 */
export const useSubgraphNavigationStore = defineStore(
  'subgraphNavigation',
  () => {
    const workflowStore = useWorkflowStore()
    const canvasStore = useCanvasStore()
    const router = useRouter()
    const routeHash = useRouteHash()

    /** The currently opened subgraph. */
    const activeSubgraph = shallowRef<Subgraph>()

    /** The stack of subgraph IDs from the root graph to the currently opened subgraph. */
    const idStack = ref<string[]>([])

    /** LRU cache for viewport states. Key: `workflowPath:graphId` */
    const viewportCache = new QuickLRU<string, DragAndScaleState>({
      maxSize: VIEWPORT_CACHE_MAX_SIZE
    })

    /** Get the ID of the root graph for the currently active workflow. */
    const getCurrentRootGraphId = () => {
      const canvas = canvasStore.getCanvas()
      return canvas.graph?.rootGraph?.id ?? 'root'
    }

    /**
     * Set by saveCurrentViewport() (called from beforeLoadNewGraph) to
     * prevent onNavigated from re-saving a stale viewport during the
     * workflow switch transition. Uses setTimeout instead of rAF so the
     * flag resets even when the tab is backgrounded.
     */
    let isWorkflowSwitching = false
    // ── Helpers ──────────────────────────────────────────────────────

    /** Build a workflow-scoped cache key. */
    function buildCacheKey(
      graphId: string,
      workflowRef?: { path?: string } | null
    ): string {
      const wf = workflowRef ?? workflowStore.activeWorkflow
      const prefix = wf?.path ?? ''
      return `${prefix}:${graphId}`
    }

    /** ID of the graph currently shown on the canvas. */
    function getActiveGraphId(): string {
      const canvas = canvasStore.getCanvas()
      return canvas?.subgraph?.id ?? getCurrentRootGraphId()
    }

    // ── Navigation stack ─────────────────────────────────────────────

    /**
     * A stack representing subgraph navigation history from the root graph to
     * the current opened subgraph.
     */
    const navigationStack = computed(() =>
      idStack.value
        .map((id) => app.rootGraph.subgraphs.get(id))
        .filter(isNonNullish)
    )

    /**
     * Restore the navigation stack from a list of subgraph IDs.
     * @see exportState
     */
    const restoreState = (subgraphIds: string[]) => {
      idStack.value.length = 0
      for (const id of subgraphIds) idStack.value.push(id)
    }

    /**
     * Export the navigation stack as a list of subgraph IDs.
     * @see restoreState
     */
    const exportState = () => [...idStack.value]

    // ── Viewport save / restore ──────────────────────────────────────

    /** Get 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 for a graph. */
    function saveViewport(graphId: string, workflowRef?: object | null): void {
      const viewport = getCurrentViewport()
      if (!viewport) return
      viewportCache.set(buildCacheKey(graphId, workflowRef), viewport)
    }

    /** Apply a viewport state to the canvas. */
    function applyViewport(viewport: DragAndScaleState): void {
      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)
    }

    function restoreViewport(graphId: string): void {
      const canvas = app.canvas
      if (!canvas) return

      const expectedKey = buildCacheKey(graphId)
      const viewport = viewportCache.get(expectedKey)
      if (viewport) {
        applyViewport(viewport)
        return
      }

      // First visit — fit to content so subgraph nodes are visible
      requestAnimationFrame(() => {
        if (getActiveGraphId() !== graphId) return
        if (!canvas.graph?.nodes?.length) return
        useLitegraphService().fitView()
        // fitView changes scale/offset, so re-sync slot positions for
        // collapsed nodes whose DOM-relative measurement is now stale.
        requestAnimationFrame(() => {
          if (getActiveGraphId() !== graphId) return
          requestSlotLayoutSyncForAllNodes()
        })
      })
    }

    // ── Navigation handler ───────────────────────────────────────────

    function onNavigated(
      subgraph: Subgraph | undefined,
      prevSubgraph: Subgraph | undefined
    ): void {
      // During a workflow switch, beforeLoadNewGraph already saved the
      // outgoing viewport — skip the save here to avoid caching stale
      // canvas state from the transition.
      if (!isWorkflowSwitching) {
        if (prevSubgraph) {
          saveViewport(prevSubgraph.id)
        } else if (!prevSubgraph && subgraph) {
          saveViewport(getCurrentRootGraphId())
        }
      }

      const isInRootGraph = !subgraph
      if (isInRootGraph) {
        idStack.value.length = 0
        restoreViewport(getCurrentRootGraphId())
        return
      }

      const path = findSubgraphPathById(subgraph.rootGraph, subgraph.id)
      const isInReachableSubgraph = !!path
      if (isInReachableSubgraph) {
        idStack.value = [...path]
      } else {
        idStack.value = [subgraph.id]
      }

      restoreViewport(subgraph.id)
    }

    // ── Watchers ─────────────────────────────────────────────────────

    // Sync flush ensures we capture the outgoing viewport before any other
    // watchers or DOM updates from the same state change mutate the canvas.
    watch(
      () => workflowStore.activeSubgraph,
      (newValue, oldValue) => {
        onNavigated(newValue, oldValue)
      },
      { flush: 'sync' }
    )

    // Counter so nested/overlapping async navigations don't release
    // suppression early; gates both the canvasStore.currentGraph watcher
    // (updateHash) and the routeHash watcher to prevent re-entrant
    // navigateToHash calls during router.replace().
    let blockNavDepth = 0
    let initialLoad = true

    async function withNavBlocked<T>(op: () => Promise<T>): Promise<T> {
      blockNavDepth++
      try {
        return await op()
      } finally {
        blockNavDepth--
      }
    }

    function ensureCanvasOnRoot() {
      const root = app.rootGraph
      const canvas = canvasStore.getCanvas()
      if (!root || !canvas) return
      if (canvas.graph?.id !== root.id) canvas.setGraph(root)
    }

    async function redirectToRoot(reason: string) {
      const root = app.rootGraph
      console.warn(`[subgraphNavigation] ${reason}; redirecting to root graph`)
      try {
        await withNavBlocked(() => router.replace('#' + root.id))
      } catch (err) {
        if (
          !isNavigationFailure(err, NavigationFailureType.duplicated) &&
          !isNavigationFailure(err, NavigationFailureType.cancelled)
        ) {
          console.warn(
            '[subgraphNavigation] router.replace rejected during recovery',
            err
          )
        }
      } finally {
        ensureCanvasOnRoot()
      }
    }

    async function navigateToHash(newHash: string) {
      const root = app.rootGraph
      const locatorId = newHash?.slice(1) || root.id
      const canvas = canvasStore.getCanvas()

      const isRoot = locatorId === root.id
      const targetGraph = isRoot
        ? root
        : isUuidShapedSubgraphId(locatorId)
          ? root.subgraphs.get(locatorId)
          : undefined
      if (targetGraph) {
        if (canvas.graph?.id === targetGraph.id) return
        return canvas.setGraph(targetGraph)
      }

      //Search all open workflows
      for (const workflow of workflowStore.openWorkflows) {
        const { activeState } = workflow
        if (!activeState) continue
        const subgraphs = activeState.definitions?.subgraphs ?? []
        for (const graph of [activeState, ...subgraphs]) {
          if (graph.id !== locatorId) continue
          // This will trigger a navigation, which can break forward history.
          // After openWorkflow resolves, app.rootGraph has been swapped, so we
          // intentionally re-read app.rootGraph below instead of using the
          // `root` captured at function entry.
          try {
            await withNavBlocked(() =>
              useWorkflowService().openWorkflow(workflow)
            )
          } catch (err) {
            console.warn(
              '[subgraphNavigation] openWorkflow rejected during recovery',
              err
            )
            return redirectToRoot('workflow load failed')
          }
          const loadedGraph =
            app.rootGraph.id === locatorId
              ? app.rootGraph
              : app.rootGraph.subgraphs.get(locatorId)
          if (!loadedGraph) {
            return redirectToRoot('subgraph not found after workflow load')
          }
          if (canvas.graph?.id === loadedGraph.id) return
          return canvas.setGraph(loadedGraph)
        }
      }

      await redirectToRoot(`subgraph not found: ${locatorId}`)
    }

    async function safeRouterCall(op: () => Promise<unknown>, label: string) {
      try {
        await op()
      } catch (err) {
        if (!isNavigationFailure(err, NavigationFailureType.duplicated)) {
          console.warn(`[subgraphNavigation] ${label} rejected`, err)
        }
      }
    }

    async function updateHash() {
      if (blockNavDepth > 0) return
      if (initialLoad) {
        initialLoad = false
        if (!routeHash.value) return
        await navigateToHash(routeHash.value)
        const graph = canvasStore.getCanvas().graph
        if (isSubgraph(graph)) workflowStore.activeSubgraph = graph
        return
      }

      const newId = canvasStore.getCanvas().graph?.id ?? ''
      if (!routeHash.value) {
        await safeRouterCall(
          () => router.replace('#' + app.rootGraph.id),
          'router.replace'
        )
      }
      const currentId = routeHash.value?.slice(1)
      if (!newId || newId === currentId) return

      await safeRouterCall(() => router.push('#' + newId), 'router.push')
    }
    watch(() => canvasStore.currentGraph, updateHash)
    watch(routeHash, () => {
      if (blockNavDepth > 0) return
      void navigateToHash(String(routeHash.value))
    })

    /** Save the current viewport for the active graph/workflow. Called by
     *  workflowService.beforeLoadNewGraph() before the canvas is overwritten. */
    function saveCurrentViewport(): void {
      saveViewport(getActiveGraphId())
      isWorkflowSwitching = true
      setTimeout(() => {
        isWorkflowSwitching = false
      }, 0)
    }

    return {
      activeSubgraph,
      navigationStack,
      restoreState,
      exportState,
      saveViewport,
      restoreViewport,
      saveCurrentViewport,
      updateHash,
      /** @internal Exposed for test assertions only. */
      viewportCache
    }
  }
)