diff --git a/src/stores/executionStore.ts b/src/stores/executionStore.ts index 9d12e20220..c2e37cd732 100644 --- a/src/stores/executionStore.ts +++ b/src/stores/executionStore.ts @@ -1,3 +1,4 @@ +import type { LGraph, Subgraph } from '@comfyorg/litegraph' import { defineStore } from 'pinia' import { computed, ref } from 'vue' @@ -22,7 +23,7 @@ import type { import { api } from '@/scripts/api' import { useCanvasStore } from './graphStore' -import { ComfyWorkflow } from './workflowStore' +import { ComfyWorkflow, useWorkflowStore } from './workflowStore' export interface QueuedPrompt { /** @@ -37,6 +38,7 @@ export interface QueuedPrompt { } export const useExecutionStore = defineStore('execution', () => { + const workflowStore = useWorkflowStore() const canvasStore = useCanvasStore() const clientId = ref(null) @@ -61,6 +63,59 @@ export const useExecutionStore = defineStore('execution', () => { ) }) + 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(null) const executingNodeProgress = computed(() => @@ -167,12 +222,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 } } @@ -193,19 +252,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) { const { nodeId, text } = e.detail if (!text || !nodeId) return - const node = canvasStore.getCanvas().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) { - const { node_id, component, props = {} } = e.detail - const node = canvasStore.getCanvas().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') { diff --git a/src/stores/workflowStore.ts b/src/stores/workflowStore.ts index 24c1f01065..6b0044e9a2 100644 --- a/src/stores/workflowStore.ts +++ b/src/stores/workflowStore.ts @@ -1,4 +1,4 @@ -import type { Subgraph } from '@comfyorg/litegraph' +import type { LGraph, Subgraph } from '@comfyorg/litegraph' import _ from 'lodash' import { defineStore } from 'pinia' import { type Raw, computed, markRaw, ref, shallowRef, watch } from 'vue' @@ -162,6 +162,7 @@ export interface WorkflowStore { activeSubgraph: Subgraph | undefined /** Updates the {@link subgraphNamePath} and {@link isSubgraphActive} values. */ updateActiveGraph: () => void + executionIdToCurrentId: (id: string) => any } export const useWorkflowStore = defineStore('workflow', () => { @@ -442,11 +443,46 @@ export const useWorkflowStore = defineStore('workflow', () => { isSubgraphActive.value = isSubgraph(subgraph) } + const subgraphNodeIdToSubgraph = (id: string, graph: LGraph | Subgraph) => { + const node = graph.getNodeById(id) + if (node?.isSubgraphNode()) return node.subgraph + } + + 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 + 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) } }