From e4110dd254e79f496a1aa67e7a784a83a3d2668e Mon Sep 17 00:00:00 2001 From: Csongor Czezar Date: Sat, 20 Dec 2025 15:40:21 -0800 Subject: [PATCH] fix: QPO progress bar now shows node name in subgraphs --- src/stores/executionStore.ts | 24 +-- .../executionStore.executingNode.test.ts | 157 ++++++++++++++++++ 2 files changed, 165 insertions(+), 16 deletions(-) create mode 100644 tests-ui/tests/store/executionStore.executingNode.test.ts diff --git a/src/stores/executionStore.ts b/src/stores/executionStore.ts index b9f89c29a..57b52bf67 100644 --- a/src/stores/executionStore.ts +++ b/src/stores/executionStore.ts @@ -7,11 +7,7 @@ import { isCloud } from '@/platform/distribution/types' import { useTelemetry } from '@/platform/telemetry' import type { ComfyWorkflow } from '@/platform/workflow/management/stores/workflowStore' import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore' -import type { - ComfyNode, - ComfyWorkflowJSON, - NodeId -} from '@/platform/workflow/validation/schemas/workflowSchema' +import type { NodeId } from '@/platform/workflow/validation/schemas/workflowSchema' import { useCanvasStore } from '@/renderer/core/canvas/canvasStore' import type { ExecutedWsMessage, @@ -194,20 +190,16 @@ export const useExecutionStore = defineStore('execution', () => { ) // For backward compatibility - returns the primary executing node - const executingNode = computed(() => { + const executingNode = computed(() => { if (!executingNodeId.value) return null - const workflow: ComfyWorkflow | undefined = activePrompt.value?.workflow - if (!workflow) return null - - const canvasState: ComfyWorkflowJSON | null = - workflow.changeTracker?.activeState ?? null - if (!canvasState) return null - - return ( - canvasState.nodes.find((n) => String(n.id) === executingNodeId.value) ?? - null + // Use getNodeByExecutionId to find nodes even in subgraphs + // executingNodeId may be a simple ID like "123" or a hierarchical ID like "123:456" + const node = getNodeByExecutionId( + app.rootGraph, + String(executingNodeId.value) ) + return node ?? null }) // This is the progress of the currently executing node (for backward compatibility) diff --git a/tests-ui/tests/store/executionStore.executingNode.test.ts b/tests-ui/tests/store/executionStore.executingNode.test.ts new file mode 100644 index 000000000..f242eb2df --- /dev/null +++ b/tests-ui/tests/store/executionStore.executingNode.test.ts @@ -0,0 +1,157 @@ +import { createPinia, setActivePinia } from 'pinia' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import type { LGraphNode } from '@/lib/litegraph/src/litegraph' +import { useExecutionStore } from '@/stores/executionStore' + +// Create mock functions +const mockNodeExecutionIdToNodeLocatorId = vi.fn() +const mockNodeIdToNodeLocatorId = vi.fn() +const mockNodeLocatorIdToNodeExecutionId = vi.fn() + +// Create a mocked graph that we can manipulate +let mockRootGraph: any + +// Mock the app import with proper implementation +vi.mock('@/scripts/app', () => ({ + app: { + get rootGraph() { + return mockRootGraph + }, + revokePreviews: vi.fn(), + nodePreviewImages: {} + } +})) + +// Mock the workflowStore +vi.mock('@/platform/workflow/management/stores/workflowStore', async () => { + const { ComfyWorkflow } = await vi.importActual< + typeof import('@/platform/workflow/management/stores/workflowStore') + >('@/platform/workflow/management/stores/workflowStore') + return { + ComfyWorkflow, + useWorkflowStore: vi.fn(() => ({ + nodeExecutionIdToNodeLocatorId: mockNodeExecutionIdToNodeLocatorId, + nodeIdToNodeLocatorId: mockNodeIdToNodeLocatorId, + nodeLocatorIdToNodeExecutionId: mockNodeLocatorIdToNodeExecutionId + })) + } +}) + +vi.mock('@/composables/node/useNodeProgressText', () => ({ + useNodeProgressText: () => ({ + showTextPreview: vi.fn() + }) +})) + +describe('useExecutionStore - executingNode with subgraphs', () => { + beforeEach(() => { + setActivePinia(createPinia()) + vi.clearAllMocks() + + // Reset the mock root graph + mockRootGraph = { + getNodeById: vi.fn(), + nodes: [] + } + }) + + it('should find executing node in root graph', () => { + const mockNode = { + id: '123', + title: 'Test Node', + type: 'TestNode' + } as LGraphNode + + mockRootGraph.getNodeById = vi.fn((id) => { + return id === '123' ? mockNode : null + }) + + const store = useExecutionStore() + + // Simulate node execution starting + store.nodeProgressStates = { + '123': { + state: 'running', + value: 0, + max: 100, + display_node_id: '123', + prompt_id: 'test-prompt', + node_id: '123' + } + } + + expect(store.executingNode).toBe(mockNode) + }) + + it('should find executing node in subgraph using execution ID', () => { + const mockNodeInSubgraph = { + id: '789', + title: 'Nested Node', + type: 'NestedNode' + } as LGraphNode + + const mockSubgraphNode = { + id: '456', + title: 'Node In Subgraph', + type: 'SubgraphNode', + isSubgraphNode: () => true, + subgraph: { + id: 'sub-uuid', + getNodeById: vi.fn((id) => { + return id === '789' ? mockNodeInSubgraph : null + }), + _nodes: [] + } + } as unknown as LGraphNode + + // Mock the graph traversal + mockRootGraph.getNodeById = vi.fn((id) => { + return id === '456' ? mockSubgraphNode : null + }) + + const store = useExecutionStore() + + // Simulate node execution in subgraph with hierarchical execution ID "456:789" + store.nodeProgressStates = { + '456:789': { + state: 'running', + value: 0, + max: 100, + display_node_id: '456:789', + prompt_id: 'test-prompt', + node_id: '456:789' + } + } + + // The executingNode should resolve to the nested node + expect(store.executingNode).toBe(mockNodeInSubgraph) + }) + + it('should return null when no node is executing', () => { + const store = useExecutionStore() + + store.nodeProgressStates = {} + + expect(store.executingNode).toBeNull() + }) + + it('should return null when executing node cannot be found', () => { + mockRootGraph.getNodeById = vi.fn(() => null) + + const store = useExecutionStore() + + store.nodeProgressStates = { + '999': { + state: 'running', + value: 0, + max: 100, + display_node_id: '999', + prompt_id: 'test-prompt', + node_id: '999' + } + } + + expect(store.executingNode).toBeNull() + }) +})