fix: Handle multiple outputs in ComponentToolkit and clean edge logic (#12268)

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]

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 = {

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);

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: