Rework TaskInput and TaskOutput logic

Author: camielvs

src/components/shared/ReactFlow/FlowCanvas/IONode/IONode.tsx

@@ -27,12 +27,7 @@ interface IONodeProps {
 }
 
 const IONode = ({ type, data, selected = false }: IONodeProps) => {
-  const {
-    getInputNodeId,
-    getOutputNodeId,
-    getTaskInputNodeId,
-    getTaskOutputNodeId,
-  } = useNodeManager();
+  const { getInputNodeId, getOutputNodeId } = useNodeManager();
   const { graphSpec, componentSpec } = useComponentSpec();
   const { setContent, clearContent } = useContextPanel();
 
@@ -63,13 +58,12 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
     [componentSpec.outputs, spec.name],
   );
 
-  const nodeId = isInput
-    ? getInputNodeId(inputNameToInputId(spec.name))
-    : getOutputNodeId(outputNameToOutputId(spec.name));
+  const inputId = inputNameToInputId(spec.name);
+  const outputId = outputNameToOutputId(spec.name);
 
-  const nodeHandleId = isInput
-    ? getTaskOutputNodeId(inputNameToInputId(spec.name))
-    : getTaskInputNodeId(outputNameToOutputId(spec.name));
+  const nodeId = isInput ? getInputNodeId(inputId) : getOutputNodeId(outputId);
+
+  const nodeHandleId = `${nodeId}_handle`;
 
   const handleHandleClick = useCallback(() => {
     if (ENABLE_DEBUG_MODE) {

src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/Handles.tsx

@@ -16,10 +16,6 @@ import { cn } from "@/lib/utils";
 import { useTaskNode } from "@/providers/TaskNodeProvider";
 import type { InputSpec, OutputSpec } from "@/utils/componentSpec";
 import { ENABLE_DEBUG_MODE } from "@/utils/constants";
-import {
-  inputNameToInputId,
-  outputNameToOutputId,
-} from "@/utils/nodes/nodeIdUtils";
 
 type InputHandleProps = {
   input: InputSpec;
@@ -39,7 +35,7 @@ export const InputHandle = ({
   onHandleSelectionChange,
 }: InputHandleProps) => {
   const { getTaskInputNodeId } = useNodeManager();
-  const { nodeId, state, name } = useTaskNode();
+  const { taskId, nodeId, state, name } = useTaskNode();
 
   const fromHandle = useConnection((connection) => connection.fromHandle?.id);
   const toHandle = useConnection((connection) => connection.toHandle?.id);
@@ -51,7 +47,7 @@ export const InputHandle = ({
   const [selected, setSelected] = useState(false);
   const [active, setActive] = useState(false);
 
-  const handleId = getTaskInputNodeId(inputNameToInputId(input.name));
+  const handleId = getTaskInputNodeId(taskId, input.name);
 
   const missing = invalid ? "bg-red-700!" : "bg-gray-500!";
   const hasValue = value !== undefined && value !== null;
@@ -235,7 +231,7 @@ export const OutputHandle = ({
   onHandleSelectionChange,
 }: OutputHandleProps) => {
   const { getTaskOutputNodeId } = useNodeManager();
-  const { nodeId, state, name } = useTaskNode();
+  const { taskId, nodeId, state, name } = useTaskNode();
 
   const fromHandle = useConnection((connection) => connection.fromHandle?.id);
   const toHandle = useConnection((connection) => connection.toHandle?.id);
@@ -247,7 +243,7 @@ export const OutputHandle = ({
   const [selected, setSelected] = useState(false);
   const [active, setActive] = useState(false);
 
-  const handleId = getTaskOutputNodeId(outputNameToOutputId(output.name));
+  const handleId = getTaskOutputNodeId(taskId, output.name);
   const hasValue = value !== undefined && value !== "" && value !== null;
 
   const handleHandleClick = useCallback(

src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/TaskNodeInputs.tsx

@@ -11,7 +11,6 @@ import { useTaskNode } from "@/providers/TaskNodeProvider";
 import { inputsWithInvalidArguments } from "@/services/componentService";
 import type { InputSpec } from "@/utils/componentSpec";
 import { ComponentSearchFilter } from "@/utils/constants";
-import { inputNameToInputId } from "@/utils/nodes/nodeIdUtils";
 import { checkArtifactMatchesSearchFilters } from "@/utils/searchUtils";
 
 import { InputHandle } from "./Handles";
@@ -28,8 +27,8 @@ export function TaskNodeInputs({
   expanded,
   onBackgroundClick,
 }: TaskNodeInputsProps) {
-  const { getInputNodeId } = useNodeManager();
-  const { inputs, taskSpec, state, select } = useTaskNode();
+  const { getTaskInputNodeId } = useNodeManager();
+  const { taskId, inputs, taskSpec, state, select } = useTaskNode();
   const { graphSpec } = useComponentSpec();
   const {
     highlightSearchFilter,
@@ -147,7 +146,7 @@ export function TaskNodeInputs({
     }
 
     const input = inputs.find(
-      (i) => getInputNodeId(inputNameToInputId(i.name)) === fromHandle?.id,
+      (i) => getTaskInputNodeId(taskId, i.name) === fromHandle?.id,
     );
 
     if (!input) return;

src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/TaskNodeOutputs.tsx

@@ -8,7 +8,6 @@ import { isValidFilterRequest } from "@/providers/ComponentLibraryProvider/types
 import { useTaskNode } from "@/providers/TaskNodeProvider";
 import type { OutputSpec } from "@/utils/componentSpec";
 import { ComponentSearchFilter } from "@/utils/constants";
-import { outputNameToOutputId } from "@/utils/nodes/nodeIdUtils";
 import { checkArtifactMatchesSearchFilters } from "@/utils/searchUtils";
 
 import { OutputHandle } from "./Handles";
@@ -25,7 +24,7 @@ export function TaskNodeOutputs({
   onBackgroundClick,
 }: TaskNodeOutputsProps) {
   const { getTaskOutputNodeId } = useNodeManager();
-  const { nodeId, outputs, state, select } = useTaskNode();
+  const { taskId, nodeId, outputs, state, select } = useTaskNode();
   const {
     highlightSearchFilter,
     resetSearchFilter,
@@ -42,8 +41,7 @@ export function TaskNodeOutputs({
     edges.some(
       (edge) =>
         edge.source === nodeId &&
-        edge.sourceHandle ===
-          getTaskOutputNodeId(outputNameToOutputId(output.name)),
+        edge.sourceHandle === getTaskOutputNodeId(taskId, output.name),
     ),
   );
 
@@ -141,8 +139,7 @@ export function TaskNodeOutputs({
     }
 
     const output = outputs.find(
-      (o) =>
-        getTaskOutputNodeId(outputNameToOutputId(o.name)) === fromHandle?.id,
+      (o) => getTaskOutputNodeId(taskId, o.name) === fromHandle?.id,
     );
 
     if (!output) return;

src/components/shared/ReactFlow/FlowCanvas/utils/addAndConnectNode.ts

@@ -10,8 +10,8 @@ import {
 } from "@/utils/componentSpec";
 import { DEFAULT_NODE_DIMENSIONS } from "@/utils/constants";
 import {
-  inputNameToInputId,
-  outputNameToOutputId,
+  inputIdToInputName,
+  outputIdToOutputName,
 } from "@/utils/nodes/nodeIdUtils";
 
 import addTask from "./addTask";
@@ -44,8 +44,58 @@ export function addAndConnectNode({
 
   const oldGraphSpec = componentSpec.implementation.graph;
 
-  const fromHandleId = fromHandle?.id;
-  const fromHandleType = fromHandleId?.startsWith("input") ? "input" : "output";
+  if (!fromHandle?.id) {
+    return componentSpec;
+  }
+
+  const isTaskHandle = nodeManager.isManaged(fromHandle.id);
+  let fromHandleType: "input" | "output";
+  let fromHandleName: string | undefined;
+  let fromTaskId: string | undefined;
+
+  if (isTaskHandle) {
+    // Handle is managed by NodeManager (task handle)
+    const fromHandleInfo = nodeManager.getHandleInfo(fromHandle.id);
+    const fromNodeType = nodeManager.getNodeType(fromHandle.id);
+
+    if (!fromHandleInfo || !fromNodeType) {
+      return componentSpec;
+    }
+
+    fromHandleType = fromNodeType === "taskInput" ? "input" : "output";
+    fromHandleName = fromHandleInfo.handleName;
+    fromTaskId = fromHandleInfo.taskId;
+  } else {
+    // Simple IO node handle - get info from the source node, not the handle
+    const fromNodeId = fromHandle.nodeId;
+    const fromNodeType = nodeManager.getNodeType(fromNodeId);
+
+    if (!fromNodeType) {
+      return componentSpec;
+    }
+
+    if (fromNodeType === "input") {
+      fromHandleType = "output";
+      const inputId = nodeManager.getTaskId(fromNodeId);
+      if (inputId) {
+        fromHandleName = inputIdToInputName(inputId);
+        fromTaskId = inputId;
+      }
+    } else if (fromNodeType === "output") {
+      fromHandleType = "input";
+      const outputId = nodeManager.getTaskId(fromNodeId);
+      if (outputId) {
+        fromHandleName = outputIdToOutputName(outputId);
+        fromTaskId = outputId;
+      }
+    } else {
+      return componentSpec;
+    }
+  }
+
+  if (!fromTaskId || !fromHandleName) {
+    return componentSpec;
+  }
 
   const adjustedPosition =
     fromHandleType === "input"
@@ -77,22 +127,17 @@ export function addAndConnectNode({
   const newNodeId = nodeManager.getNodeId(newTaskId, "task");
 
   // 3. Determine the connection data type and find the first matching handle on the new node
-  if (!fromHandle) {
-    return newComponentSpec;
+  let fromComponentSpec: ComponentSpec | undefined;
+
+  if (isTaskHandle) {
+    // Get spec from task
+    const fromTaskSpec = graphSpec.tasks[fromTaskId];
+    fromComponentSpec = fromTaskSpec?.componentRef.spec;
+  } else {
+    // For IO nodes, get spec from component spec
+    fromComponentSpec = componentSpec;
   }
 
-  const fromTaskId = nodeManager.getTaskId(fromHandle.nodeId);
-  if (!fromTaskId) {
-    return newComponentSpec;
-  }
-
-  const fromTaskSpec = graphSpec.tasks[fromTaskId];
-  const fromComponentSpec = fromTaskSpec?.componentRef.spec;
-
-  const fromNodeId = fromHandle.nodeId;
-
-  const fromHandleName = fromHandleId?.replace(`${fromHandleType}_`, "");
-
   let connectionType: TypeSpecType | undefined;
   if (fromHandleType === "input") {
     connectionType = fromComponentSpec?.inputs?.find(
@@ -106,7 +151,6 @@ export function addAndConnectNode({
 
   // Find the first matching handle on the new node
   const toHandleType = fromHandleType === "input" ? "output" : "input";
-
   let targetHandleId: string | undefined;
 
   if (toHandleType === "input") {
@@ -117,8 +161,11 @@ export function addAndConnectNode({
       return newComponentSpec;
     }
 
-    const inputId = inputNameToInputId(handleName);
-    targetHandleId = nodeManager.getNodeId(inputId, "taskInput");
+    targetHandleId = nodeManager.getTaskHandleNodeId(
+      newTaskId,
+      handleName,
+      "taskInput",
+    );
   } else if (toHandleType === "output") {
     const handleName = componentRef.spec?.outputs?.find(
       (io) => io.type === connectionType,
@@ -127,14 +174,21 @@ export function addAndConnectNode({
       return newComponentSpec;
     }
 
-    const outputId = outputNameToOutputId(handleName);
-    targetHandleId = nodeManager.getNodeId(outputId, "taskOutput");
+    targetHandleId = nodeManager.getTaskHandleNodeId(
+      newTaskId,
+      handleName,
+      "taskOutput",
+    );
   }
 
   // 4. Build a Connection object and use handleConnection to add the edge
-  if (fromNodeId && fromHandleId && targetHandleId) {
+  if (targetHandleId) {
+    const fromNodeId = fromHandle.nodeId;
+    const fromHandleId = fromHandle.id;
+
     const isReversedConnection =
       fromHandleType === "input" && toHandleType === "output";
+
     const connection: Connection = isReversedConnection
       ? // Drawing from an input handle to a new output handle
         {