Add Subgraphs (#3905)

Co-authored-by: github-actions <github-actions@github.com>

src/stores/domWidgetStore.ts

@@ -2,7 +2,7 @@
  * Stores all DOM widgets that are used in the canvas.
  */
 import { defineStore } from 'pinia'
-import { type Raw, markRaw, ref } from 'vue'
+import { type Raw, computed, markRaw, ref } from 'vue'
 
 import type { PositionConfig } from '@/composables/element/useAbsolutePosition'
 import type { BaseDOMWidget } from '@/scripts/domWidget'
@@ -13,11 +13,20 @@ export interface DomWidgetState extends PositionConfig {
   visible: boolean
   readonly: boolean
   zIndex: number
+  /** If the widget belongs to the current graph/subgraph. */
+  active: boolean
 }
 
 export const useDomWidgetStore = defineStore('domWidget', () => {
   const widgetStates = ref<Map<string, DomWidgetState>>(new Map())
 
+  const activeWidgetStates = computed(() =>
+    [...widgetStates.value.values()].filter((state) => state.active)
+  )
+  const inactiveWidgetStates = computed(() =>
+    [...widgetStates.value.values()].filter((state) => !state.active)
+  )
+
   // Register a widget with the store
   const registerWidget = <V extends object | string>(
     widget: BaseDOMWidget<V>
@@ -28,7 +37,8 @@ export const useDomWidgetStore = defineStore('domWidget', () => {
       readonly: false,
       zIndex: 0,
       pos: [0, 0],
-      size: [0, 0]
+      size: [0, 0],
+      active: true
     })
   }
 
@@ -37,9 +47,28 @@ export const useDomWidgetStore = defineStore('domWidget', () => {
     widgetStates.value.delete(widgetId)
   }
 
+  const activateWidget = (widgetId: string) => {
+    const state = widgetStates.value.get(widgetId)
+    if (state) state.active = true
+  }
+
+  const deactivateWidget = (widgetId: string) => {
+    const state = widgetStates.value.get(widgetId)
+    if (state) state.active = false
+  }
+
+  const clear = () => {
+    widgetStates.value.clear()
+  }
+
   return {
     widgetStates,
+    activeWidgetStates,
+    inactiveWidgetStates,
     registerWidget,
-    unregisterWidget
+    unregisterWidget,
+    activateWidget,
+    deactivateWidget,
+    clear
   }
 })

src/stores/executionStore.ts

@@ -1,3 +1,4 @@
+import type { LGraph, Subgraph } from '@comfyorg/litegraph'
 import { defineStore } from 'pinia'
 import { computed, ref } from 'vue'
 
@@ -20,9 +21,9 @@ import type {
   NodeId
 } from '@/schemas/comfyWorkflowSchema'
 import { api } from '@/scripts/api'
-import { app } from '@/scripts/app'
 
-import { ComfyWorkflow } from './workflowStore'
+import { useCanvasStore } from './graphStore'
+import { ComfyWorkflow, useWorkflowStore } from './workflowStore'
 
 export interface QueuedPrompt {
   /**
@@ -37,6 +38,9 @@ export interface QueuedPrompt {
 }
 
 export const useExecutionStore = defineStore('execution', () => {
+  const workflowStore = useWorkflowStore()
+  const canvasStore = useCanvasStore()
+
   const clientId = ref<string | null>(null)
   const activePromptId = ref<string | null>(null)
   const queuedPrompts = ref<Record<NodeId, QueuedPrompt>>({})
@@ -54,12 +58,64 @@ export const useExecutionStore = defineStore('execution', () => {
     if (!canvasState) return null
 
     return (
-      canvasState.nodes.find(
-        (n: ComfyNode) => String(n.id) === executingNodeId.value
-      ) ?? null
+      canvasState.nodes.find((n) => String(n.id) === executingNodeId.value) ??
+      null
     )
   })
 
+  const subgraphNodeIdToSubgraph = (id: string, graph: LGraph | Subgraph) => {
+    const node = graph.getNodeById(id)
+    if (node?.isSubgraphNode()) return node.subgraph
+  }
+
+  /**
+   * Recursively get the subgraph objects for the given subgraph instance IDs
+   * @param currentGraph The current graph
+   * @param subgraphNodeIds The instance IDs
+   * @param subgraphs The subgraphs
+   * @returns The subgraphs that correspond to each of the instance IDs.
+   */
+  const getSubgraphsFromInstanceIds = (
+    currentGraph: LGraph | Subgraph,
+    subgraphNodeIds: string[],
+    subgraphs: Subgraph[] = []
+  ): Subgraph[] => {
+    // Last segment is the node portion; nothing to do.
+    if (subgraphNodeIds.length === 1) return subgraphs
+
+    const currentPart = subgraphNodeIds.shift()
+    if (currentPart === undefined) return subgraphs
+
+    const subgraph = subgraphNodeIdToSubgraph(currentPart, currentGraph)
+    if (!subgraph) throw new Error(`Subgraph not found: ${currentPart}`)
+
+    subgraphs.push(subgraph)
+    return getSubgraphsFromInstanceIds(subgraph, subgraphNodeIds, subgraphs)
+  }
+
+  const executionIdToCurrentId = (id: string) => {
+    const subgraph = workflowStore.activeSubgraph
+
+    // Short-circuit: ID belongs to the parent workflow / no active subgraph
+    if (!id.includes(':')) {
+      return !subgraph ? id : undefined
+    } else if (!subgraph) {
+      return
+    }
+
+    // Parse the hierarchical ID (e.g., "123:456:789")
+    const subgraphNodeIds = id.split(':')
+
+    // If the last subgraph is the active subgraph, return the node ID
+    const subgraphs = getSubgraphsFromInstanceIds(
+      subgraph.rootGraph,
+      subgraphNodeIds
+    )
+    if (subgraphs.at(-1) === subgraph) {
+      return subgraphNodeIds.at(-1)
+    }
+  }
+
   // This is the progress of the currently executing node, if any
   const _executingNodeProgress = ref<ProgressWsMessage | null>(null)
   const executingNodeProgress = computed(() =>
@@ -132,7 +188,7 @@ export const useExecutionStore = defineStore('execution', () => {
     activePrompt.value.nodes[e.detail.node] = true
   }
 
-  function handleExecuting(e: CustomEvent<NodeId | null>) {
+  function handleExecuting(e: CustomEvent<NodeId | null>): void {
     // Clear the current node progress when a new node starts executing
     _executingNodeProgress.value = null
 
@@ -142,12 +198,16 @@ export const useExecutionStore = defineStore('execution', () => {
       // Seems sometimes nodes that are cached fire executing but not executed
       activePrompt.value.nodes[executingNodeId.value] = true
     }
-    executingNodeId.value = e.detail
-    if (executingNodeId.value === null) {
-      if (activePromptId.value) {
-        delete queuedPrompts.value[activePromptId.value]
+    if (typeof e.detail === 'string') {
+      executingNodeId.value = executionIdToCurrentId(e.detail) ?? null
+    } else {
+      executingNodeId.value = e.detail
+      if (executingNodeId.value === null) {
+        if (activePromptId.value) {
+          delete queuedPrompts.value[activePromptId.value]
+        }
+        activePromptId.value = null
       }
-      activePromptId.value = null
     }
   }
 
@@ -168,19 +228,31 @@ export const useExecutionStore = defineStore('execution', () => {
     lastExecutionError.value = e.detail
   }
 
+  function getNodeIdIfExecuting(nodeId: string | number) {
+    const nodeIdStr = String(nodeId)
+    return nodeIdStr.includes(':')
+      ? workflowStore.executionIdToCurrentId(nodeIdStr)
+      : nodeIdStr
+  }
+
   function handleProgressText(e: CustomEvent<ProgressTextWsMessage>) {
     const { nodeId, text } = e.detail
     if (!text || !nodeId) return
 
-    const node = app.graph.getNodeById(nodeId)
+    // Handle hierarchical node IDs for subgraphs
+    const currentId = getNodeIdIfExecuting(nodeId)
+    const node = canvasStore.getCanvas().graph?.getNodeById(currentId)
     if (!node) return
 
     useNodeProgressText().showTextPreview(node, text)
   }
 
   function handleDisplayComponent(e: CustomEvent<DisplayComponentWsMessage>) {
-    const { node_id, component, props = {} } = e.detail
-    const node = app.graph.getNodeById(node_id)
+    const { node_id: nodeId, component, props = {} } = e.detail
+
+    // Handle hierarchical node IDs for subgraphs
+    const currentId = getNodeIdIfExecuting(nodeId)
+    const node = canvasStore.getCanvas().graph?.getNodeById(currentId)
     if (!node) return
 
     if (component === 'ChatHistoryWidget') {

src/stores/graphStore.ts

@@ -3,7 +3,7 @@ import type { Positionable } from '@comfyorg/litegraph/dist/interfaces'
 import { defineStore } from 'pinia'
 import { type Raw, computed, markRaw, ref, shallowRef } from 'vue'
 
-import { isLGraphGroup, isLGraphNode } from '@/utils/litegraphUtil'
+import { isLGraphGroup, isLGraphNode, isReroute } from '@/utils/litegraphUtil'
 
 export const useTitleEditorStore = defineStore('titleEditor', () => {
   const titleEditorTarget = shallowRef<LGraphNode | LGraphGroup | null>(null)
@@ -31,6 +31,7 @@ export const useCanvasStore = defineStore('canvas', () => {
 
   const nodeSelected = computed(() => selectedItems.value.some(isLGraphNode))
   const groupSelected = computed(() => selectedItems.value.some(isLGraphGroup))
+  const rerouteSelected = computed(() => selectedItems.value.some(isReroute))
 
   const getCanvas = () => {
     if (!canvas.value) throw new Error('getCanvas: canvas is null')
@@ -42,6 +43,7 @@ export const useCanvasStore = defineStore('canvas', () => {
     selectedItems,
     nodeSelected,
     groupSelected,
+    rerouteSelected,
     updateSelectedItems,
     getCanvas
   }

src/stores/nodeDefStore.ts

@@ -306,8 +306,7 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
   }
   function fromLGraphNode(node: LGraphNode): ComfyNodeDefImpl | null {
     // Frontend-only nodes don't have nodeDef
-    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-    // @ts-ignore Optional chaining used in index
+    // @ts-expect-error Optional chaining used in index
     return nodeDefsByName.value[node.constructor?.nodeData?.name] ?? null
   }
 

src/stores/subgraphNavigationStore.ts

@@ -0,0 +1,87 @@
+import type { Subgraph } from '@comfyorg/litegraph'
+import { defineStore } from 'pinia'
+import { computed, shallowReactive, shallowRef, watch } from 'vue'
+
+import { app } from '@/scripts/app'
+import { isNonNullish } from '@/utils/typeGuardUtil'
+
+import { useWorkflowStore } from './workflowStore'
+
+/**
+ * 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()
+
+    /** The currently opened subgraph. */
+    const activeSubgraph = shallowRef<Subgraph>()
+
+    /** The stack of subgraph IDs from the root graph to the currently opened subgraph. */
+    const idStack = shallowReactive<string[]>([])
+
+    /**
+     * A stack representing subgraph navigation history from the root graph to
+     * the current opened subgraph.
+     */
+    const navigationStack = computed(() =>
+      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,
+      () => (idStack.length = 0)
+    )
+
+    // Update navigation stack when opened subgraph changes
+    watch(
+      () => workflowStore.activeSubgraph,
+      (subgraph) => {
+        // Navigated back to the root graph
+        if (!subgraph) {
+          idStack.length = 0
+          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)
+        }
+      }
+    )
+
+    return {
+      activeSubgraph,
+      navigationStack,
+      restoreState,
+      exportState
+    }
+  }
+)

src/stores/workflowStore.ts

@@ -1,6 +1,7 @@
+import type { LGraph, Subgraph } from '@comfyorg/litegraph'
 import _ from 'lodash'
 import { defineStore } from 'pinia'
-import { computed, markRaw, ref, watch } from 'vue'
+import { type Raw, computed, markRaw, ref, shallowRef, watch } from 'vue'
 
 import { ComfyWorkflowJSON } from '@/schemas/comfyWorkflowSchema'
 import { api } from '@/scripts/api'
@@ -156,12 +157,12 @@ export interface WorkflowStore {
   syncWorkflows: (dir?: string) => Promise<void>
   reorderWorkflows: (from: number, to: number) => void
 
-  /** An ordered list of all parent subgraphs, ending with the current subgraph. */
-  subgraphNamePath: string[]
   /** `true` if any subgraph is currently being viewed. */
   isSubgraphActive: boolean
+  activeSubgraph: Subgraph | undefined
   /** Updates the {@link subgraphNamePath} and {@link isSubgraphActive} values. */
   updateActiveGraph: () => void
+  executionIdToCurrentId: (id: string) => any
 }
 
 export const useWorkflowStore = defineStore('workflow', () => {
@@ -427,24 +428,61 @@ export const useWorkflowStore = defineStore('workflow', () => {
     }
   }
 
-  /** @see WorkflowStore.subgraphNamePath */
-  const subgraphNamePath = ref<string[]>([])
   /** @see WorkflowStore.isSubgraphActive */
   const isSubgraphActive = ref(false)
 
+  /** @see WorkflowStore.activeSubgraph */
+  const activeSubgraph = shallowRef<Raw<Subgraph>>()
+
   /** @see WorkflowStore.updateActiveGraph */
   const updateActiveGraph = () => {
+    const subgraph = comfyApp.canvas?.subgraph
+    activeSubgraph.value = subgraph ? markRaw(subgraph) : undefined
     if (!comfyApp.canvas) return
 
-    const { subgraph } = comfyApp.canvas
     isSubgraphActive.value = isSubgraph(subgraph)
+  }
 
-    if (subgraph) {
-      const [, ...pathFromRoot] = subgraph.pathToRootGraph
+  const subgraphNodeIdToSubgraph = (id: string, graph: LGraph | Subgraph) => {
+    const node = graph.getNodeById(id)
+    if (node?.isSubgraphNode()) return node.subgraph
+  }
 
-      subgraphNamePath.value = pathFromRoot.map((graph) => graph.name)
-    } else {
-      subgraphNamePath.value = []
+  const getSubgraphsFromInstanceIds = (
+    currentGraph: LGraph | Subgraph,
+    subgraphNodeIds: string[],
+    subgraphs: Subgraph[] = []
+  ): Subgraph[] => {
+    const currentPart = subgraphNodeIds.shift()
+    if (currentPart === undefined) return subgraphs
+
+    const subgraph = subgraphNodeIdToSubgraph(currentPart, currentGraph)
+    if (subgraph === undefined) throw new Error('Subgraph not found')
+
+    subgraphs.push(subgraph)
+    return getSubgraphsFromInstanceIds(subgraph, subgraphNodeIds, subgraphs)
+  }
+
+  const executionIdToCurrentId = (id: string) => {
+    const subgraph = activeSubgraph.value
+
+    // Short-circuit: ID belongs to the parent workflow / no active subgraph
+    if (!id.includes(':')) {
+      return !subgraph ? id : undefined
+    } else if (!subgraph) {
+      return
+    }
+
+    // Parse the hierarchical ID (e.g., "123:456:789")
+    const subgraphNodeIds = id.split(':')
+
+    // Start from the root graph
+    const { graph } = comfyApp
+
+    // If the last subgraph is the active subgraph, return the node ID
+    const subgraphs = getSubgraphsFromInstanceIds(graph, subgraphNodeIds)
+    if (subgraphs.at(-1) === subgraph) {
+      return subgraphNodeIds.at(-1)
     }
   }
 
@@ -473,9 +511,10 @@ export const useWorkflowStore = defineStore('workflow', () => {
     getWorkflowByPath,
     syncWorkflows,
 
-    subgraphNamePath,
     isSubgraphActive,
-    updateActiveGraph
+    activeSubgraph,
+    updateActiveGraph,
+    executionIdToCurrentId
   }
 }) satisfies () => WorkflowStore