From 70eab1b2a02d70a46a56e8c993ccc694e38ac2d5 Mon Sep 17 00:00:00 2001 From: Csaba Tuncsik Date: Fri, 25 Jul 2025 16:12:21 +0200 Subject: [PATCH] fix(core): Optimize connection type lookups (#17585) --- .../src/composables/useWorkflowHelpers.ts | 2 +- packages/workflow/src/workflow-data-proxy.ts | 4 +- packages/workflow/src/workflow.ts | 75 ++- packages/workflow/test/workflow.test.ts | 473 ++++++++++++++++++ 4 files changed, 508 insertions(+), 46 deletions(-) diff --git a/packages/frontend/editor-ui/src/composables/useWorkflowHelpers.ts b/packages/frontend/editor-ui/src/composables/useWorkflowHelpers.ts index 9bdbc3daf4..4a3ace5411 100644 --- a/packages/frontend/editor-ui/src/composables/useWorkflowHelpers.ts +++ b/packages/frontend/editor-ui/src/composables/useWorkflowHelpers.ts @@ -169,7 +169,7 @@ function resolveParameterImpl( let contextNode = activeNode; if (activeNode) { - contextNode = workflow.getParentMainInputNode(activeNode); + contextNode = workflow.getParentMainInputNode(activeNode) ?? null; } const workflowRunData = executionData?.data?.resultData.runData ?? null; diff --git a/packages/workflow/src/workflow-data-proxy.ts b/packages/workflow/src/workflow-data-proxy.ts index 554dace0d6..74eb79ae5c 100644 --- a/packages/workflow/src/workflow-data-proxy.ts +++ b/packages/workflow/src/workflow-data-proxy.ts @@ -1095,7 +1095,7 @@ export class WorkflowDataProxy { let contextNode = that.contextNodeName; if (activeNode) { const parentMainInputNode = that.workflow.getParentMainInputNode(activeNode); - contextNode = parentMainInputNode.name ?? contextNode; + contextNode = parentMainInputNode?.name ?? contextNode; } if (!that.workflow.hasPath(nodeName, contextNode)) { @@ -1147,7 +1147,7 @@ export class WorkflowDataProxy { let contextNode = that.contextNodeName; if (activeNode) { const parentMainInputNode = that.workflow.getParentMainInputNode(activeNode); - contextNode = parentMainInputNode.name ?? contextNode; + contextNode = parentMainInputNode?.name ?? contextNode; } // Use bidirectional path checking to handle AI/tool nodes properly diff --git a/packages/workflow/src/workflow.ts b/packages/workflow/src/workflow.ts index 2b00afe648..a8f6ea893a 100644 --- a/packages/workflow/src/workflow.ts +++ b/packages/workflow/src/workflow.ts @@ -36,7 +36,6 @@ import type { INodeConnection, IObservableObject, NodeParameterValueType, - INodeOutputConfiguration, NodeConnectionType, } from './interfaces'; import { NodeConnectionTypes } from './interfaces'; @@ -677,40 +676,28 @@ export class Workflow { return returnConns; } - getParentMainInputNode(node: INode): INode { - if (node) { - const nodeType = this.nodeTypes.getByNameAndVersion(node.type, node.typeVersion); - const outputs = NodeHelpers.getNodeOutputs(this, node, nodeType.description); + getParentMainInputNode(node: INode | null | undefined): INode | null | undefined { + if (!node) return node; - if ( - outputs.find( - (output) => - ((output as INodeOutputConfiguration)?.type ?? output) !== NodeConnectionTypes.Main, - ) - ) { - // Get the first node which is connected to a non-main output - const nonMainNodesConnected = outputs?.reduce((acc, outputName) => { - const parentNodes = this.getChildNodes( - node.name, - (outputName as INodeOutputConfiguration)?.type ?? outputName, - ); - if (parentNodes.length > 0) { - acc.push(...parentNodes); + const nodeConnections = this.connectionsBySourceNode[node.name]; + if (!nodeConnections) return node; + + // Get non-main connection types + const nonMainConnectionTypes = Object.keys(nodeConnections).filter( + (type) => type !== NodeConnectionTypes.Main, + ); + + for (const connectionType of nonMainConnectionTypes) { + const connections = nodeConnections[connectionType] ?? []; + for (const connectionGroup of connections) { + for (const connection of connectionGroup ?? []) { + if (connection?.node) { + const returnNode = this.getNode(connection.node); + if (!returnNode) { + throw new ApplicationError(`Node "${connection.node}" not found`); + } + return this.getParentMainInputNode(returnNode); } - return acc; - }, [] as string[]); - - if (nonMainNodesConnected.length) { - const returnNode = this.getNode(nonMainNodesConnected[0]); - if (returnNode === null) { - // This should theoretically never happen as the node is connected - // but who knows and it makes TS happy - throw new ApplicationError(`Node "${nonMainNodesConnected[0]}" not found`); - } - - // The chain of non-main nodes is potentially not finished yet so - // keep on going - return this.getParentMainInputNode(returnNode); } } } @@ -933,6 +920,17 @@ export class Workflow { hasPath(fromNodeName: string, toNodeName: string, maxDepth = 50): boolean { if (fromNodeName === toNodeName) return true; + // Get connection types that actually exist in this workflow + // We need both source and destination connection types for bidirectional search + const connectionTypes = new Set(); + for (const nodeConnections of Object.values(this.connectionsBySourceNode).concat( + Object.values(this.connectionsByDestinationNode), + )) { + for (const type of Object.keys(nodeConnections)) { + connectionTypes.add(type as NodeConnectionType); + } + } + const visited = new Set(); const queue: Array<{ nodeName: string; depth: number }> = [ { nodeName: fromNodeName, depth: 0 }, @@ -947,16 +945,7 @@ export class Workflow { visited.add(nodeName); - // Check all connection types for this node - const allConnectionTypes = [ - NodeConnectionTypes.Main, - NodeConnectionTypes.AiTool, - NodeConnectionTypes.AiMemory, - NodeConnectionTypes.AiDocument, - NodeConnectionTypes.AiVectorStore, - ]; - - for (const connectionType of allConnectionTypes) { + for (const connectionType of connectionTypes) { // Get children (forward direction) const children = this.getChildNodes(nodeName, connectionType); for (const childName of children) { diff --git a/packages/workflow/test/workflow.test.ts b/packages/workflow/test/workflow.test.ts index fd0c8e8c1a..02cf9873b1 100644 --- a/packages/workflow/test/workflow.test.ts +++ b/packages/workflow/test/workflow.test.ts @@ -2362,6 +2362,479 @@ describe('Workflow', () => { const result = WORKFLOW_WITH_LOOPS.getParentMainInputNode(set1Node); expect(result).toBe(set1Node); }); + + describe('nodes with only main outputs', () => { + test('should return the same node when it only has main outputs', () => { + const nodes: INode[] = [ + { + id: '1', + name: 'SimpleNode', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + { + id: '2', + name: 'TargetNode', + type: 'test.set', + typeVersion: 1, + position: [200, 100], + parameters: {}, + }, + ]; + + const connections = { + SimpleNode: { + [NodeConnectionTypes.Main]: [ + [{ node: 'TargetNode', type: NodeConnectionTypes.Main, index: 0 }], + ], + }, + }; + + const workflow = new Workflow({ + id: 'test', + nodes, + connections, + active: false, + nodeTypes, + }); + + const simpleNode = workflow.getNode('SimpleNode')!; + const result = workflow.getParentMainInputNode(simpleNode); + + expect(result).toBe(simpleNode); + expect(result!.name).toBe('SimpleNode'); + }); + + test('should return the same node when it has no connections', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'IsolatedNode', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + ], + connections: {}, + active: false, + nodeTypes, + }); + + const isolatedNode = workflow.getNode('IsolatedNode')!; + const result = workflow.getParentMainInputNode(isolatedNode); + + expect(result).toBe(isolatedNode); + expect(result!.name).toBe('IsolatedNode'); + }); + }); + + describe('nodes with non-main outputs (AI/Tool connections)', () => { + test('should follow AI tool connection to find main input node', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'ToolNode', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + { + id: '2', + name: 'AgentNode', + type: 'test.set', + typeVersion: 1, + position: [200, 100], + parameters: {}, + }, + ], + connections: { + ToolNode: { + [NodeConnectionTypes.AiTool]: [ + [{ node: 'AgentNode', type: NodeConnectionTypes.AiTool, index: 0 }], + ], + }, + }, + active: false, + nodeTypes, + }); + + const toolNode = workflow.getNode('ToolNode')!; + const result = workflow.getParentMainInputNode(toolNode); + + expect(result!.name).toBe('AgentNode'); + }); + + test('should follow AI memory connection to find main input node', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'MemoryNode', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + { + id: '2', + name: 'ChatNode', + type: 'test.set', + typeVersion: 1, + position: [200, 100], + parameters: {}, + }, + ], + connections: { + MemoryNode: { + [NodeConnectionTypes.AiMemory]: [ + [{ node: 'ChatNode', type: NodeConnectionTypes.AiMemory, index: 0 }], + ], + }, + }, + active: false, + nodeTypes, + }); + + const memoryNode = workflow.getNode('MemoryNode')!; + const result = workflow.getParentMainInputNode(memoryNode); + + expect(result!.name).toBe('ChatNode'); + }); + + test('should handle mixed main and non-main outputs', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'MixedNode', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + { + id: '2', + name: 'MainTarget', + type: 'test.set', + typeVersion: 1, + position: [200, 100], + parameters: {}, + }, + { + id: '3', + name: 'ToolTarget', + type: 'test.set', + typeVersion: 1, + position: [200, 200], + parameters: {}, + }, + ], + connections: { + MixedNode: { + [NodeConnectionTypes.Main]: [ + [{ node: 'MainTarget', type: NodeConnectionTypes.Main, index: 0 }], + ], + [NodeConnectionTypes.AiTool]: [ + [{ node: 'ToolTarget', type: NodeConnectionTypes.AiTool, index: 0 }], + ], + }, + }, + active: false, + nodeTypes, + }); + + const mixedNode = workflow.getNode('MixedNode')!; + const result = workflow.getParentMainInputNode(mixedNode); + + // Should follow the first non-main connection (AiTool) + expect(result!.name).toBe('ToolTarget'); + }); + }); + + describe('chain traversal scenarios', () => { + test('should follow a chain of AI connections until reaching main input node', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'StartTool', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + { + id: '2', + name: 'MiddleTool', + type: 'test.set', + typeVersion: 1, + position: [200, 100], + parameters: {}, + }, + { + id: '3', + name: 'FinalAgent', + type: 'test.set', + typeVersion: 1, + position: [300, 100], + parameters: {}, + }, + ], + connections: { + StartTool: { + [NodeConnectionTypes.AiTool]: [ + [{ node: 'MiddleTool', type: NodeConnectionTypes.AiTool, index: 0 }], + ], + }, + MiddleTool: { + [NodeConnectionTypes.AiTool]: [ + [{ node: 'FinalAgent', type: NodeConnectionTypes.AiTool, index: 0 }], + ], + }, + }, + active: false, + nodeTypes, + }); + + const startTool = workflow.getNode('StartTool')!; + const result = workflow.getParentMainInputNode(startTool); + + expect(result!.name).toBe('FinalAgent'); + }); + + test('should handle chain that ends with a node having only main outputs', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'ToolNode', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + { + id: '2', + name: 'IntermediateNode', + type: 'test.set', + typeVersion: 1, + position: [200, 100], + parameters: {}, + }, + { + id: '3', + name: 'EndNode', + type: 'test.set', + typeVersion: 1, + position: [300, 100], + parameters: {}, + }, + ], + connections: { + ToolNode: { + [NodeConnectionTypes.AiTool]: [ + [{ node: 'IntermediateNode', type: NodeConnectionTypes.AiTool, index: 0 }], + ], + }, + IntermediateNode: { + [NodeConnectionTypes.Main]: [ + [{ node: 'EndNode', type: NodeConnectionTypes.Main, index: 0 }], + ], + }, + }, + active: false, + nodeTypes, + }); + + const toolNode = workflow.getNode('ToolNode')!; + const result = workflow.getParentMainInputNode(toolNode); + + expect(result!.name).toBe('IntermediateNode'); + }); + + test('should handle complex multi-branch AI connections', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'MultiTool', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + { + id: '2', + name: 'Agent1', + type: 'test.set', + typeVersion: 1, + position: [200, 50], + parameters: {}, + }, + { + id: '3', + name: 'Agent2', + type: 'test.set', + typeVersion: 1, + position: [200, 150], + parameters: {}, + }, + ], + connections: { + MultiTool: { + [NodeConnectionTypes.AiTool]: [ + [ + { node: 'Agent1', type: NodeConnectionTypes.AiTool, index: 0 }, + { node: 'Agent2', type: NodeConnectionTypes.AiTool, index: 0 }, + ], + ], + }, + }, + active: false, + nodeTypes, + }); + + const multiTool = workflow.getNode('MultiTool')!; + const result = workflow.getParentMainInputNode(multiTool); + + // Should follow the first connection in the array + expect(result!.name).toBe('Agent1'); + }); + }); + + describe('edge cases', () => { + test('should handle null node input', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [], + connections: {}, + active: false, + nodeTypes, + }); + + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + const result = workflow.getParentMainInputNode(null as any); + expect(result).toBeNull(); + }); + + test('should handle undefined node input', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [], + connections: {}, + active: false, + nodeTypes, + }); + + // eslint-disable-next-line @typescript-eslint/no-unsafe-argument + const result = workflow.getParentMainInputNode(undefined as any); + expect(result).toBeUndefined(); + }); + + test('should throw error when connected node does not exist in workflow', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'ToolNode', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + ], + connections: { + ToolNode: { + [NodeConnectionTypes.AiTool]: [ + [{ node: 'NonExistentNode', type: NodeConnectionTypes.AiTool, index: 0 }], + ], + }, + }, + active: false, + nodeTypes, + }); + + const toolNode = workflow.getNode('ToolNode')!; + + expect(() => { + workflow.getParentMainInputNode(toolNode); + }).toThrow('Node "NonExistentNode" not found'); + }); + + test('should handle empty connection arrays', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'EmptyConnectionNode', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + ], + connections: { + EmptyConnectionNode: { + [NodeConnectionTypes.AiTool]: [ + [], // Empty connection array + ], + }, + }, + active: false, + nodeTypes, + }); + + const emptyConnectionNode = workflow.getNode('EmptyConnectionNode')!; + const result = workflow.getParentMainInputNode(emptyConnectionNode); + + expect(result).toBe(emptyConnectionNode); + }); + + test('should handle null connections in connection array', () => { + const workflow = new Workflow({ + id: 'test', + nodes: [ + { + id: '1', + name: 'NullConnectionNode', + type: 'test.set', + typeVersion: 1, + position: [100, 100], + parameters: {}, + }, + ], + connections: { + NullConnectionNode: { + [NodeConnectionTypes.AiTool]: [ + [{ node: '', type: NodeConnectionTypes.AiTool, index: 0 }], // Connection with empty node name + ], + }, + }, + active: false, + nodeTypes, + }); + + const nullConnectionNode = workflow.getNode('NullConnectionNode')!; + const result = workflow.getParentMainInputNode(nullConnectionNode); + + expect(result).toBe(nullConnectionNode); + }); + }); }); describe('getNodeConnectionIndexes', () => {