From 72e08191b0c393262ef8207c6f0b7bf26f5d8de1 Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa Date: Wed, 25 Mar 2026 08:48:47 -0300 Subject: [PATCH] fix: Handle multiple outputs in ComponentToolkit and clean edge logic (#12268) --- ...test_component_toolkit_multiple_outputs.py | 80 ++++++++++++ .../src/utils/__tests__/cleanEdges.test.ts | 121 +++++++++++++++++- src/frontend/src/utils/reactflowUtils.ts | 77 +++++++++-- src/lfx/src/lfx/base/tools/component_tool.py | 8 +- 4 files changed, 264 insertions(+), 22 deletions(-) create mode 100644 src/backend/tests/unit/base/tools/test_component_toolkit_multiple_outputs.py diff --git a/src/backend/tests/unit/base/tools/test_component_toolkit_multiple_outputs.py b/src/backend/tests/unit/base/tools/test_component_toolkit_multiple_outputs.py new file mode 100644 index 0000000000..49e87233e3 --- /dev/null +++ b/src/backend/tests/unit/base/tools/test_component_toolkit_multiple_outputs.py @@ -0,0 +1,80 @@ +"""Tests for ComponentToolkit handling of components with multiple outputs. + +Bug: When a component has 2+ outputs (each with tool_mode=True) and +tool_name/tool_description are provided, ComponentToolkit.get_tools() raises +ValueError instead of disambiguating tool names. + +GIVEN: A component with 2 outputs, both with tool_mode=True (default) +WHEN: get_tools() is called with tool_name and tool_description +THEN: ValueError: "When passing a tool name or description, there must be + only one tool, but 2 tools were found." +EXPECTED: Tools are created with prefixed names to disambiguate them. +""" + +from lfx.base.tools.component_tool import ComponentToolkit +from lfx.custom.custom_component.component import Component +from lfx.io import MessageTextInput, Output +from lfx.schema.message import Message + + +class MultiOutputComponent(Component): + """Minimal component with two outputs to reproduce the bug.""" + + display_name = "Multi Output Component" + description = "A component with two separate outputs." + name = "MultiOutputComponent" + + inputs = [ + MessageTextInput( + name="input_value", + display_name="Input", + tool_mode=True, + ), + ] + + outputs = [ + Output( + display_name="Progress", + name="progress_output", + method="get_progress", + ), + Output( + display_name="Result", + name="result_output", + method="get_result", + ), + ] + + async def get_progress(self) -> Message: + return await Message.create(text="progress") + + async def get_result(self) -> Message: + return await Message.create(text="result") + + +def test_should_create_tools_with_prefixed_names_when_component_has_multiple_outputs(): + """Bug fix: get_tools() must handle multiple outputs with tool_name/description. + + Instead of raising ValueError, it should prefix each tool's name + with the provided tool_name to disambiguate them. + """ + # Arrange + component = MultiOutputComponent() + toolkit = ComponentToolkit(component=component) + tool_name = "my_agent" + tool_description = "An agent with split outputs" + + # Act + tools = toolkit.get_tools( + tool_name=tool_name, + tool_description=tool_description, + ) + + # Assert — must create 2 tools with prefixed names, not raise ValueError + assert len(tools) == 2 + tool_names = {tool.name for tool in tools} + assert "my_agent_get_progress" in tool_names + assert "my_agent_get_result" in tool_names + for tool in tools: + assert tool_description in tool.description + assert tool.tags == [tool.name] diff --git a/src/frontend/src/utils/__tests__/cleanEdges.test.ts b/src/frontend/src/utils/__tests__/cleanEdges.test.ts index 3860b2499e..1c6d77ffbc 100644 --- a/src/frontend/src/utils/__tests__/cleanEdges.test.ts +++ b/src/frontend/src/utils/__tests__/cleanEdges.test.ts @@ -17,8 +17,25 @@ const scapeJSONParse = (str: string) => { } }; +const customStringify = (obj: any): string => { + if (typeof obj === "undefined") return "null"; + if (obj === null || typeof obj !== "object") { + if (obj instanceof Date) return `"${obj.toISOString()}"`; + return JSON.stringify(obj); + } + if (Array.isArray(obj)) { + const arrayItems = obj.map((item) => customStringify(item)).join(","); + return `[${arrayItems}]`; + } + const keys = Object.keys(obj).sort(); + const keyValuePairs = keys.map( + (key) => `"${key}":${customStringify(obj[key])}`, + ); + return `{${keyValuePairs.join(",")}}`; +}; + const scapedJSONStringfy = (obj: unknown) => { - return JSON.stringify(obj).replace(/"/g, "œ"); + return customStringify(obj).replace(/"/g, "œ"); }; // Define minimal types needed for testing @@ -229,11 +246,17 @@ function cleanEdges(nodes: AllNodeType[], edges: EdgeType[]) { const name = parsedSourceHandle.name; if (sourceNode.type === "genericNode") { - const output = - sourceNode.data.node.outputs?.find( - (o) => o.name === sourceNode.data.selected_output, - ) ?? - sourceNode.data.node.outputs?.find( + // For components with group_outputs, each output has its own handle, + // so we must NOT use selected_output (which would match the wrong output). + const hasGroupOutputs = sourceNode.data.node.outputs?.some( + (o) => o.group_outputs, + ); + const outputBySelectedOutput = hasGroupOutputs + ? undefined + : sourceNode.data.node.outputs?.find( + (o) => o.name === sourceNode.data.selected_output, + ); + const outputByFallback = sourceNode.data.node.outputs?.find( (o) => (o.selected || (sourceNode.data.node.outputs?.filter( @@ -241,6 +264,7 @@ function cleanEdges(nodes: AllNodeType[], edges: EdgeType[]) { )?.length ?? 0) <= 1) && o.name === name, ); + const output = outputBySelectedOutput ?? outputByFallback; if (output) { const outputTypes = @@ -567,6 +591,91 @@ describe("cleanEdges", () => { }); }); + describe("group_outputs edge preservation", () => { + it("should preserve edge from non-selected output when component has group_outputs", () => { + // Bug: When a component has group_outputs: true and selected_output is set to one output, + // cleanEdges incorrectly uses selected_output to look up ALL edges, causing edges + // from non-selected outputs to be removed (name mismatch). + const sourceNode: AllNodeType = { + id: "AgentSplit-1", + type: "genericNode", + data: { + id: "AgentSplit-1", + type: "AgentWithSplitOutputs", + selected_output: "progress_output", + node: { + display_name: "Agent with Split Outputs", + template: {}, + outputs: [ + { + name: "progress_output", + types: ["Message"], + selected: "Message", + group_outputs: true, + }, + { + name: "result_output", + types: ["Message"], + selected: "Message", + group_outputs: true, + }, + ], + }, + }, + }; + + const targetNode: AllNodeType = { + id: "TypeConvert-1", + type: "genericNode", + data: { + id: "TypeConvert-1", + type: "TypeConverterComponent", + node: { + display_name: "Type Convert", + template: { + input_value: { + type: "other", + input_types: ["Message"], + }, + }, + outputs: [], + }, + }, + }; + + // Edge connects result_output (NOT the selected_output) to Type Convert + const sourceHandleStr = scapedJSONStringfy({ + dataType: "AgentWithSplitOutputs", + id: "AgentSplit-1", + name: "result_output", + output_types: ["Message"], + }); + + const targetHandleStr = scapedJSONStringfy({ + fieldName: "input_value", + id: "TypeConvert-1", + inputTypes: ["Message"], + type: "other", + }); + + const edge: EdgeType = { + id: "edge-result-to-convert", + source: "AgentSplit-1", + target: "TypeConvert-1", + sourceHandle: sourceHandleStr, + targetHandle: targetHandleStr, + }; + + const result = cleanEdges([sourceNode, targetNode], [edge]); + + // Edge MUST be preserved — the bug was that selected_output: "progress_output" + // caused cleanEdges to look up "progress_output" instead of "result_output", + // resulting in a name mismatch and edge removal + expect(result.edges.length).toBe(1); + expect(result.brokenEdges.length).toBe(0); + }); + }); + describe("Basic edge validation", () => { it("should remove edge when source node does not exist", () => { const targetNodeBasic: AllNodeType = { diff --git a/src/frontend/src/utils/reactflowUtils.ts b/src/frontend/src/utils/reactflowUtils.ts index 5e09d65bcd..ec1a85a834 100644 --- a/src/frontend/src/utils/reactflowUtils.ts +++ b/src/frontend/src/utils/reactflowUtils.ts @@ -198,14 +198,41 @@ export function cleanEdges(nodes: AllNodeType[], edges: EdgeType[]) { // Backward compatibility: old flows may have Data/DataFrame types that need to match JSON/Table const expectedTargetHandle = scapedJSONStringfy(id); + const targetHandlesMatchResult = handlesMatch( + expectedTargetHandle, + targetHandle, + ); if ( - (!handlesMatch(expectedTargetHandle, targetHandle) || + (!targetHandlesMatchResult || (targetNode.data.node?.tool_mode && isToolMode) || isAdvanced) && !isLoopInput ) { newEdges = newEdges.filter((e) => e.id !== edge.id); brokenEdges.push(generateAlertObject(sourceNode, targetNode, edge)); + } else if ( + targetHandlesMatchResult && + expectedTargetHandle !== targetHandle + ) { + // Handles match via migration but IDs differ — update edge to use current types + // so React Flow can find the DOM handle + const edgeInNewEdges = newEdges.find((e) => e.id === edge.id); + if (edgeInNewEdges) { + edgeInNewEdges.targetHandle = expectedTargetHandle; + if (edgeInNewEdges.data) { + edgeInNewEdges.data.targetHandle = scapeJSONParse( + expectedTargetHandle, + ); + } + // Update edge ID to reflect new handles + edgeInNewEdges.id = + "reactflow__edge-" + + edgeInNewEdges.source + + (edgeInNewEdges.sourceHandle ?? "") + + "-" + + edgeInNewEdges.target + + expectedTargetHandle; + } } } if (sourceHandle) { @@ -213,10 +240,18 @@ export function cleanEdges(nodes: AllNodeType[], edges: EdgeType[]) { const name = parsedSourceHandle.name; if (sourceNode.type == "genericNode") { - const output = - sourceNode.data.node!.outputs?.find( - (output) => output.name === sourceNode.data.selected_output, - ) ?? + // For components with group_outputs, each output has its own handle, + // so we must NOT use selected_output (which would match the wrong output). + // Only use selected_output for single-output dropdown components. + const hasGroupOutputs = sourceNode.data.node!.outputs?.some( + (o) => o.group_outputs, + ); + const outputBySelectedOutput = hasGroupOutputs + ? undefined + : sourceNode.data.node!.outputs?.find( + (output) => output.name === sourceNode.data.selected_output, + ); + const outputByFallback = sourceNode.data.node!.outputs?.find( (output) => (output.selected || @@ -225,6 +260,7 @@ export function cleanEdges(nodes: AllNodeType[], edges: EdgeType[]) { )?.length ?? 0) <= 1) && output.name === name, ); + const output = outputBySelectedOutput ?? outputByFallback; if (output) { const outputTypes = @@ -241,12 +277,35 @@ export function cleanEdges(nodes: AllNodeType[], edges: EdgeType[]) { const hasAllowsLoop = output?.allows_loop === true; // Backward compatibility: old flows may have Data/DataFrame types that need to match JSON/Table const expectedSourceHandle = scapedJSONStringfy(id); - if ( - !handlesMatch(expectedSourceHandle, sourceHandle) && - !hasAllowsLoop - ) { + const sourceMatchResult = handlesMatch( + expectedSourceHandle, + sourceHandle, + ); + if (!sourceMatchResult && !hasAllowsLoop) { newEdges = newEdges.filter((e) => e.id !== edge.id); brokenEdges.push(generateAlertObject(sourceNode, targetNode, edge)); + } else if ( + sourceMatchResult && + expectedSourceHandle !== sourceHandle + ) { + // Handles match via migration but IDs differ — update edge to use current types + const edgeInNewEdges = newEdges.find((e) => e.id === edge.id); + if (edgeInNewEdges) { + edgeInNewEdges.sourceHandle = expectedSourceHandle; + if (edgeInNewEdges.data) { + edgeInNewEdges.data.sourceHandle = scapeJSONParse( + expectedSourceHandle, + ); + } + // Update edge ID to reflect new handles + edgeInNewEdges.id = + "reactflow__edge-" + + edgeInNewEdges.source + + expectedSourceHandle + + "-" + + edgeInNewEdges.target + + (edgeInNewEdges.targetHandle ?? ""); + } } } else { newEdges = newEdges.filter((e) => e.id !== edge.id); diff --git a/src/lfx/src/lfx/base/tools/component_tool.py b/src/lfx/src/lfx/base/tools/component_tool.py index 2637d6b96b..81afed2ccf 100644 --- a/src/lfx/src/lfx/base/tools/component_tool.py +++ b/src/lfx/src/lfx/base/tools/component_tool.py @@ -270,19 +270,13 @@ class ComponentToolkit: tool.name = _format_tool_name(str(tool_name)) or tool.name tool.description = tool_description or tool.description tool.tags = [tool.name] - elif flow_mode_inputs and (tool_name or tool_description): + elif (tool_name or tool_description) and (flow_mode_inputs or len(tools) > 1): for tool in tools: tool.name = _format_tool_name(str(tool_name) + "_" + str(tool.name)) or tool.name tool.description = ( str(tool_description) + " Output details: " + str(tool.description) ) or tool.description tool.tags = [tool.name] - elif tool_name or tool_description: - msg = ( - "When passing a tool name or description, there must be only one tool, " - f"but {len(tools)} tools were found." - ) - raise ValueError(msg) return tools def get_tools_metadata_dictionary(self) -> dict: