Rework TaskInput and TaskOutput logic

Author: camielvs

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/conversions";
 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/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/conversions";
 
 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
         {

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

@@ -16,48 +16,60 @@ import { setTaskArgument } from "./setTaskArgument";
 
 type NodeInfo = {
   id: string;
-  handle?: string;
   type?: NodeType;
+  handle?: {
+    taskId: string;
+    handleName: string;
+  };
 };
 
 export const handleConnection = (
   graphSpec: GraphSpec,
   connection: Connection,
   nodeManager: NodeManager,
 ) => {
-  const sourceId = nodeManager.getTaskId(connection.source);
-  const targetId = nodeManager.getTaskId(connection.target);
-
-  if (!sourceId || !targetId) {
-    console.warn("Source or Target ID is missing in the connection.");
+  const sourceId = nodeManager.getTaskId(connection.source!);
+  const sourceType = nodeManager.getNodeType(connection.source!);
+
+  const targetId = nodeManager.getTaskId(connection.target!);
+  const targetType = nodeManager.getNodeType(connection.target!);
+
+  if (!sourceId || !targetId || !sourceType || !targetType) {
+    console.error("Could not resolve node information:", {
+      sourceId,
+      sourceType,
+      targetId,
+      targetType,
+    });
     return graphSpec;
   }
 
   if (sourceId === targetId) {
-    console.warn(
-      "Source and Target IDs are the same. Self-connections are not allowed.",
-    );
+    console.warn("Cannot connect node to itself");
     return graphSpec;
   }
 
-  const sourceHandleId = connection.sourceHandle
-    ? nodeManager.getTaskId(connection.sourceHandle)
-    : undefined;
+  let sourceHandleInfo: { taskId: string; handleName: string } | undefined;
+  let targetHandleInfo: { taskId: string; handleName: string } | undefined;
+
+  if (connection.sourceHandle) {
+    sourceHandleInfo = nodeManager.getHandleInfo(connection.sourceHandle);
+  }
 
-  const targetHandleId = connection.targetHandle
-    ? nodeManager.getTaskId(connection.targetHandle)
-    : undefined;
+  if (connection.targetHandle) {
+    targetHandleInfo = nodeManager.getHandleInfo(connection.targetHandle);
+  }
 
   const source: NodeInfo = {
     id: sourceId,
-    handle: sourceHandleId,
-    type: nodeManager.getNodeType(connection.source),
+    type: sourceType,
+    handle: sourceHandleInfo,
   };
 
   const target: NodeInfo = {
     id: targetId,
-    handle: targetHandleId,
-    type: nodeManager.getNodeType(connection.target),
+    type: targetType,
+    handle: targetHandleInfo,
   };
 
   const connectionType = `${source.type}_to_${target.type}` as const;
@@ -73,7 +85,7 @@ export const handleConnection = (
       return handleTaskToGraphOutput(graphSpec, source, target);
 
     default:
-      console.warn("Unsupported connection pattern:", connectionType);
+      console.error("Unsupported connection pattern:", connectionType);
       return graphSpec;
   }
 };
@@ -83,16 +95,19 @@ const handleGraphInputToTask = (
   source: NodeInfo,
   target: NodeInfo,
 ): GraphSpec => {
-  if (!target.handle) {
-    console.warn("Handle ID is missing for target task node.");
+  if (!target.handle?.handleName) {
+    console.error(
+      "Target handle name missing for graph input to task connection",
+    );
     return graphSpec;
   }
 
-  const sourceInputName = inputIdToInputName(source.id);
-  const targetInputName = inputIdToInputName(target.handle);
+  const inputId = source.id;
+  const inputName = inputIdToInputName(inputId);
+  const targetInputName = target.handle.handleName;
 
   const graphInputArgument: GraphInputArgument = {
-    graphInput: { inputName: sourceInputName },
+    graphInput: { inputName },
   };
 
   return setTaskArgument(
@@ -108,18 +123,18 @@ const handleTaskToTask = (
   source: NodeInfo,
   target: NodeInfo,
 ): GraphSpec => {
-  if (!source.handle) {
-    console.warn("Handle ID is missing for source task node.");
+  if (!source.handle?.handleName) {
+    console.error("Source handle name missing for task to task connection");
     return graphSpec;
   }
 
-  if (!target.handle) {
-    console.warn("Handle ID is missing for target task node.");
+  if (!target.handle?.handleName) {
+    console.error("Target handle name missing for task to task connection");
     return graphSpec;
   }
 
-  const sourceOutputName = outputIdToOutputName(source.handle);
-  const targetInputName = inputIdToInputName(target.handle);
+  const sourceOutputName = source.handle.handleName;
+  const targetInputName = target.handle.handleName;
 
   const taskOutputArgument: TaskOutputArgument = {
     taskOutput: {
@@ -141,12 +156,16 @@ const handleTaskToGraphOutput = (
   source: NodeInfo,
   target: NodeInfo,
 ): GraphSpec => {
-  if (!source.handle) {
-    console.warn("Handle ID is missing for source task node.");
+  if (!source.handle?.handleName) {
+    console.error(
+      "Source handle name missing for task to graph output connection",
+    );
     return graphSpec;
   }
 
-  const sourceOutputName = outputIdToOutputName(source.handle);
+  const sourceOutputName = source.handle.handleName;
+  const outputId = target.id;
+  const outputName = outputIdToOutputName(outputId);
 
   const taskOutputArgument: TaskOutputArgument = {
     taskOutput: {
@@ -155,5 +174,5 @@ const handleTaskToGraphOutput = (
     },
   };
 
-  return setGraphOutputValue(graphSpec, target.id, taskOutputArgument);
+  return setGraphOutputValue(graphSpec, outputName, taskOutputArgument);
 };