[feat] Move partial execution to the backend and make work with subgraphs (#4624)

src/utils/executionUtil.ts

@@ -1,8 +1,7 @@
 import type {
   ExecutableLGraphNode,
   ExecutionId,
-  LGraph,
-  NodeId
+  LGraph
 } from '@comfyorg/litegraph'
 import {
   ExecutableNodeDTO,
@@ -18,31 +17,6 @@ import type {
 import { ExecutableGroupNodeDTO, isGroupNode } from './executableGroupNodeDto'
 import { compressWidgetInputSlots } from './litegraphUtil'
 
-/**
- * Recursively target node's parent nodes to the new output.
- * @param nodeId The node id to add.
- * @param oldOutput The old output.
- * @param newOutput The new output.
- * @returns The new output.
- */
-function recursiveAddNodes(
-  nodeId: NodeId,
-  oldOutput: ComfyApiWorkflow,
-  newOutput: ComfyApiWorkflow
-) {
-  const currentId = String(nodeId)
-  const currentNode = oldOutput[currentId]!
-  if (newOutput[currentId] == null) {
-    newOutput[currentId] = currentNode
-    for (const inputValue of Object.values(currentNode.inputs || [])) {
-      if (Array.isArray(inputValue)) {
-        recursiveAddNodes(inputValue[0], oldOutput, newOutput)
-      }
-    }
-  }
-  return newOutput
-}
-
 /**
  * Converts the current graph workflow for sending to the API.
  * @note Node widgets are updated before serialization to prepare queueing.
@@ -50,14 +24,13 @@ function recursiveAddNodes(
  * @param graph The graph to convert.
  * @param options The options for the conversion.
  *  - `sortNodes`: Whether to sort the nodes by execution order.
- *  - `queueNodeIds`: The output nodes to execute. Execute all output nodes if not provided.
  * @returns The workflow and node links
  */
 export const graphToPrompt = async (
   graph: LGraph,
-  options: { sortNodes?: boolean; queueNodeIds?: NodeId[] } = {}
+  options: { sortNodes?: boolean } = {}
 ): Promise<{ workflow: ComfyWorkflowJSON; output: ComfyApiWorkflow }> => {
-  const { sortNodes = false, queueNodeIds } = options
+  const { sortNodes = false } = options
 
   for (const node of graph.computeExecutionOrder(false)) {
     const innerNodes = node.getInnerNodes
@@ -104,7 +77,7 @@ export const graphToPrompt = async (
     nodeDtoMap.set(dto.id, dto)
   }
 
-  let output: ComfyApiWorkflow = {}
+  const output: ComfyApiWorkflow = {}
   // Process nodes in order of execution
   for (const node of nodeDtoMap.values()) {
     // Don't serialize muted nodes
@@ -180,14 +153,5 @@ export const graphToPrompt = async (
     }
   }
 
-  // Partial execution
-  if (queueNodeIds?.length) {
-    const newOutput = {}
-    for (const queueNodeId of queueNodeIds) {
-      recursiveAddNodes(queueNodeId, output, newOutput)
-    }
-    output = newOutput
-  }
-
   return { workflow: workflow as ComfyWorkflowJSON, output }
 }

src/utils/graphTraversalUtil.ts

@@ -1,6 +1,6 @@
 import type { LGraph, LGraphNode, Subgraph } from '@comfyorg/litegraph'
 
-import type { NodeLocatorId } from '@/types/nodeIdentification'
+import type { NodeExecutionId, NodeLocatorId } from '@/types/nodeIdentification'
 import { parseNodeLocatorId } from '@/types/nodeIdentification'
 
 import { isSubgraphIoNode } from './typeGuardUtil'
@@ -351,3 +351,106 @@ export function mapSubgraphNodes<T>(
 export function getAllNonIoNodesInSubgraph(subgraph: Subgraph): LGraphNode[] {
   return subgraph.nodes.filter((node) => !isSubgraphIoNode(node))
 }
+
+/**
+ * Performs depth-first traversal of nodes and their subgraphs.
+ * Generic visitor pattern that can be used for various node processing tasks.
+ *
+ * @param nodes - Starting nodes for traversal
+ * @param visitor - Function called for each node with its context
+ * @param expandSubgraphs - Whether to traverse into subgraph nodes (default: true)
+ */
+export function traverseNodesDepthFirst<T>(
+  nodes: LGraphNode[],
+  visitor: (node: LGraphNode, context: T) => T,
+  initialContext: T,
+  expandSubgraphs: boolean = true
+): void {
+  type StackItem = { node: LGraphNode; context: T }
+  const stack: StackItem[] = []
+
+  // Initialize stack with starting nodes
+  for (const node of nodes) {
+    stack.push({ node, context: initialContext })
+  }
+
+  // Process stack iteratively (DFS)
+  while (stack.length > 0) {
+    const { node, context } = stack.pop()!
+
+    // Visit node and get updated context for children
+    const childContext = visitor(node, context)
+
+    // If it's a subgraph and we should expand, add children to stack
+    if (expandSubgraphs && node.isSubgraphNode?.() && node.subgraph) {
+      // Process children in reverse order to maintain left-to-right DFS processing
+      // when popping from stack (LIFO). Iterate backwards to avoid array reversal.
+      const children = node.subgraph.nodes
+      for (let i = children.length - 1; i >= 0; i--) {
+        stack.push({ node: children[i], context: childContext })
+      }
+    }
+  }
+}
+
+/**
+ * Collects nodes with custom data during depth-first traversal.
+ * Generic collector that can gather any type of data per node.
+ *
+ * @param nodes - Starting nodes for traversal
+ * @param collector - Function that returns data to collect for each node
+ * @param contextBuilder - Function that builds context for child nodes
+ * @param expandSubgraphs - Whether to traverse into subgraph nodes
+ * @returns Array of collected data
+ */
+export function collectFromNodes<T, C>(
+  nodes: LGraphNode[],
+  collector: (node: LGraphNode, context: C) => T | null,
+  contextBuilder: (node: LGraphNode, parentContext: C) => C,
+  initialContext: C,
+  expandSubgraphs: boolean = true
+): T[] {
+  const results: T[] = []
+
+  traverseNodesDepthFirst(
+    nodes,
+    (node, context) => {
+      const data = collector(node, context)
+      if (data !== null) {
+        results.push(data)
+      }
+      return contextBuilder(node, context)
+    },
+    initialContext,
+    expandSubgraphs
+  )
+
+  return results
+}
+
+/**
+ * Collects execution IDs for selected nodes and all their descendants.
+ * Uses the generic DFS traversal with optimized string building.
+ *
+ * @param selectedNodes - The selected nodes to process
+ * @returns Array of execution IDs for selected nodes and all nodes within selected subgraphs
+ */
+export function getExecutionIdsForSelectedNodes(
+  selectedNodes: LGraphNode[]
+): NodeExecutionId[] {
+  return collectFromNodes(
+    selectedNodes,
+    // Collector: build execution ID for each node
+    (node, parentExecutionId: string): NodeExecutionId => {
+      const nodeId = String(node.id)
+      return parentExecutionId ? `${parentExecutionId}:${nodeId}` : nodeId
+    },
+    // Context builder: pass execution ID to children
+    (node, parentExecutionId: string) => {
+      const nodeId = String(node.id)
+      return parentExecutionId ? `${parentExecutionId}:${nodeId}` : nodeId
+    },
+    '', // Initial context: empty parent execution ID
+    true // Expand subgraphs
+  )
+}

src/utils/nodeFilterUtil.ts

@@ -0,0 +1,21 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+
+/**
+ * Checks if a node is an output node.
+ * Output nodes are nodes that have the output_node flag set in their nodeData.
+ *
+ * @param node - The node to check
+ * @returns True if the node is an output node, false otherwise
+ */
+export const isOutputNode = (node: LGraphNode) =>
+  node.constructor.nodeData?.output_node
+
+/**
+ * Filters nodes to find only output nodes.
+ * Output nodes are nodes that have the output_node flag set in their nodeData.
+ *
+ * @param nodes - Array of nodes to filter
+ * @returns Array of output nodes only
+ */
+export const filterOutputNodes = (nodes: LGraphNode[]): LGraphNode[] =>
+  nodes.filter(isOutputNode)