Fix inability to delete edges

camielvs · camielvs · commit b62b77dd3da5 · 2025-10-09T10:15:12.000-07:00
diff --git a/src/components/shared/ReactFlow/FlowCanvas/IONode/IONode.tsx b/src/components/shared/ReactFlow/FlowCanvas/IONode/IONode.tsx
@@ -28,7 +28,7 @@ interface IONodeProps {
 }
 
 const IONode = ({ type, data, selected = false }: IONodeProps) => {
-  const { getInputNodeId, getOutputNodeId } = useNodeManager();
+  const { getNodeId, getHandleNodeId } = useNodeManager();
   const { graphSpec, componentSpec } = useComponentSpec();
   const { setContent, clearContent } = useContextPanel();
 
@@ -59,11 +59,13 @@ 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 id = isInput ? inputId : outputId;
 
-  const nodeHandleId = `${nodeId}_handle`;
+  const nodeId = getNodeId(id, type);
+  const handleNodeType = isInput ? "outputHandle" : "inputHandle";
+  const nodeHandleId = getHandleNodeId(id, "handle", handleNodeType);
 
   const handleHandleClick = useCallback(() => {
     if (ENABLE_DEBUG_MODE) {
diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/Handles.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/Handles.tsx
@@ -34,7 +34,7 @@ export const InputHandle = ({
   onLabelClick,
   onHandleSelectionChange,
 }: InputHandleProps) => {
-  const { getTaskInputNodeId } = useNodeManager();
+  const { getInputHandleNodeId } = useNodeManager();
   const { taskId, nodeId, state, name } = useTaskNode();
 
   const fromHandle = useConnection((connection) => connection.fromHandle?.id);
@@ -47,7 +47,7 @@ export const InputHandle = ({
   const [selected, setSelected] = useState(false);
   const [active, setActive] = useState(false);
 
-  const handleId = getTaskInputNodeId(taskId, input.name);
+  const handleId = getInputHandleNodeId(taskId, input.name);
 
   const missing = invalid ? "bg-red-700!" : "bg-gray-500!";
   const hasValue = value !== undefined && value !== null;
@@ -230,7 +230,7 @@ export const OutputHandle = ({
   onLabelClick,
   onHandleSelectionChange,
 }: OutputHandleProps) => {
-  const { getTaskOutputNodeId } = useNodeManager();
+  const { getOutputHandleNodeId } = useNodeManager();
   const { taskId, nodeId, state, name } = useTaskNode();
 
   const fromHandle = useConnection((connection) => connection.fromHandle?.id);
@@ -243,7 +243,7 @@ export const OutputHandle = ({
   const [selected, setSelected] = useState(false);
   const [active, setActive] = useState(false);
 
-  const handleId = getTaskOutputNodeId(taskId, output.name);
+  const handleId = getOutputHandleNodeId(taskId, output.name);
   const hasValue = value !== undefined && value !== "" && value !== null;
 
   const handleHandleClick = useCallback(
diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/TaskNodeInputs.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/TaskNodeInputs.tsx
@@ -27,7 +27,7 @@ export function TaskNodeInputs({
   expanded,
   onBackgroundClick,
 }: TaskNodeInputsProps) {
-  const { getTaskInputNodeId } = useNodeManager();
+  const { getInputHandleNodeId } = useNodeManager();
   const { taskId, inputs, taskSpec, state, select } = useTaskNode();
   const { graphSpec } = useComponentSpec();
   const {
@@ -146,7 +146,7 @@ export function TaskNodeInputs({
     }
 
     const input = inputs.find(
-      (i) => getTaskInputNodeId(taskId, i.name) === fromHandle?.id,
+      (i) => getInputHandleNodeId(taskId, i.name) === fromHandle?.id,
     );
 
     if (!input) return;
diff --git a/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/TaskNodeOutputs.tsx b/src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/TaskNodeOutputs.tsx
@@ -23,7 +23,7 @@ export function TaskNodeOutputs({
   expanded,
   onBackgroundClick,
 }: TaskNodeOutputsProps) {
-  const { getTaskOutputNodeId } = useNodeManager();
+  const { getOutputHandleNodeId } = useNodeManager();
   const { taskId, nodeId, outputs, state, select } = useTaskNode();
   const {
     highlightSearchFilter,
@@ -41,7 +41,7 @@ export function TaskNodeOutputs({
     edges.some(
       (edge) =>
         edge.source === nodeId &&
-        edge.sourceHandle === getTaskOutputNodeId(taskId, output.name),
+        edge.sourceHandle === getOutputHandleNodeId(taskId, output.name),
     ),
   );
 
@@ -139,7 +139,7 @@ export function TaskNodeOutputs({
     }
 
     const output = outputs.find(
-      (o) => getTaskOutputNodeId(taskId, o.name) === fromHandle?.id,
+      (o) => getOutputHandleNodeId(taskId, o.name) === fromHandle?.id,
     );
 
     if (!output) return;
diff --git a/src/components/shared/ReactFlow/FlowCanvas/utils/addAndConnectNode.ts b/src/components/shared/ReactFlow/FlowCanvas/utils/addAndConnectNode.ts
@@ -1,6 +1,6 @@
 import type { Connection, Handle } from "@xyflow/react";
 
-import type { NodeManager } from "@/nodeManager";
+import type { NodeManager, NodeType } from "@/nodeManager";
 import {
   type ComponentReference,
   type ComponentSpec,
@@ -48,57 +48,44 @@ export function addAndConnectNode({
     return componentSpec;
   }
 
-  const isTaskHandle = nodeManager.isManaged(fromHandle.id);
-  let fromHandleType: "input" | "output";
-  let fromHandleName: string | undefined;
-  let fromTaskId: string | undefined;
+  const fromNodeId = fromHandle.nodeId;
+  const fromNodeType = nodeManager.getNodeType(fromNodeId);
+  const fromTaskId = nodeManager.getTaskId(fromNodeId);
 
-  if (isTaskHandle) {
-    // Handle is managed by NodeManager (task handle)
-    const fromHandleInfo = nodeManager.getHandleInfo(fromHandle.id);
-    const fromNodeType = nodeManager.getNodeType(fromHandle.id);
+  if (!fromTaskId) {
+    return componentSpec;
+  }
 
-    if (!fromHandleInfo || !fromNodeType) {
-      return componentSpec;
-    }
+  let fromHandleType: NodeType | undefined;
+  let fromHandleName: string | undefined;
 
-    fromHandleType = fromNodeType === "taskInput" ? "input" : "output";
-    fromHandleName = fromHandleInfo.handleName;
-    fromTaskId = fromHandleInfo.taskId;
+  if (fromNodeType === "task") {
+    const fromHandleInfo = nodeManager.getHandleInfo(fromHandle.id);
+    fromHandleName = fromHandleInfo?.handleName;
+    fromHandleType = nodeManager.getNodeType(fromHandle.id);
+  } else if (fromNodeType === "input") {
+    fromHandleType = "outputHandle";
+    fromHandleName = inputIdToInputName(fromTaskId);
+  } else if (fromNodeType === "output") {
+    fromHandleType = "inputHandle";
+    fromHandleName = outputIdToOutputName(fromTaskId);
   } 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;
-    }
+    return componentSpec;
+  }
+
+  if (!fromHandleName) {
+    return componentSpec;
   }
 
-  if (!fromTaskId || !fromHandleName) {
+  if (
+    !fromHandleType ||
+    (fromHandleType !== "inputHandle" && fromHandleType !== "outputHandle")
+  ) {
     return componentSpec;
   }
 
   const adjustedPosition =
-    fromHandleType === "input"
+    fromHandleType === "inputHandle"
       ? { ...position, x: position.x - DEFAULT_NODE_DIMENSIONS.w }
       : position;
 
@@ -129,7 +116,7 @@ export function addAndConnectNode({
   // 3. Determine the connection data type and find the first matching handle on the new node
   let fromComponentSpec: ComponentSpec | undefined;
 
-  if (isTaskHandle) {
+  if (fromNodeType === "task") {
     // Get spec from task
     const fromTaskSpec = graphSpec.tasks[fromTaskId];
     fromComponentSpec = fromTaskSpec?.componentRef.spec;
@@ -139,55 +126,48 @@ export function addAndConnectNode({
   }
 
   let connectionType: TypeSpecType | undefined;
-  if (fromHandleType === "input") {
+  if (fromHandleType === "inputHandle") {
     connectionType = fromComponentSpec?.inputs?.find(
       (io) => io.name === fromHandleName,
     )?.type;
-  } else if (fromHandleType === "output") {
+  } else if (fromHandleType === "outputHandle") {
     connectionType = fromComponentSpec?.outputs?.find(
       (io) => io.name === fromHandleName,
     )?.type;
   }
 
   // Find the first matching handle on the new node
-  const toHandleType = fromHandleType === "input" ? "output" : "input";
-  let targetHandleId: string | undefined;
-
-  if (toHandleType === "input") {
-    const handleName = componentRef.spec?.inputs?.find(
-      (io) => io.type === connectionType,
-    )?.name;
-    if (!handleName) {
-      return newComponentSpec;
-    }
-
-    targetHandleId = nodeManager.getTaskHandleNodeId(
-      newTaskId,
-      handleName,
-      "taskInput",
-    );
-  } else if (toHandleType === "output") {
-    const handleName = componentRef.spec?.outputs?.find(
-      (io) => io.type === connectionType,
-    )?.name;
-    if (!handleName) {
-      return newComponentSpec;
-    }
-
-    targetHandleId = nodeManager.getTaskHandleNodeId(
-      newTaskId,
-      handleName,
-      "taskOutput",
-    );
+  const toHandleType =
+    fromHandleType === "inputHandle" ? "outputHandle" : "inputHandle";
+
+  const inputHandleName = componentRef.spec?.inputs?.find(
+    (io) => io.type === connectionType,
+  )?.name;
+
+  const outputHandleName = componentRef.spec?.outputs?.find(
+    (io) => io.type === connectionType,
+  )?.name;
+
+  const toHandleName =
+    toHandleType === "inputHandle" ? inputHandleName : outputHandleName;
+
+  if (!toHandleName) {
+    return newComponentSpec;
   }
 
+  const targetHandleId = nodeManager.getHandleNodeId(
+    newTaskId,
+    toHandleName,
+    toHandleType,
+  );
+
   // 4. Build a Connection object and use handleConnection to add the edge
   if (targetHandleId) {
     const fromNodeId = fromHandle.nodeId;
     const fromHandleId = fromHandle.id;
 
     const isReversedConnection =
-      fromHandleType === "input" && toHandleType === "output";
+      fromHandleType === "inputHandle" && toHandleType === "outputHandle";
 
     const connection: Connection = isReversedConnection
       ? // Drawing from an input handle to a new output handle
diff --git a/src/components/shared/ReactFlow/FlowCanvas/utils/removeEdge.ts b/src/components/shared/ReactFlow/FlowCanvas/utils/removeEdge.ts
@@ -1,7 +1,10 @@
 import type { Edge } from "@xyflow/react";
 
 import type { NodeManager } from "@/nodeManager";
-import type { ComponentSpec, GraphImplementation } from "@/utils/componentSpec";
+import {
+  type ComponentSpec,
+  isGraphImplementation,
+} from "@/utils/componentSpec";
 import { outputIdToOutputName } from "@/utils/nodes/conversions";
 
 import { setGraphOutputValue } from "./setGraphOutputValue";
@@ -12,38 +15,67 @@ export const removeEdge = (
   componentSpec: ComponentSpec,
   nodeManager: NodeManager,
 ) => {
-  const graphSpec = (componentSpec.implementation as GraphImplementation)
-    ?.graph;
+  if (!isGraphImplementation(componentSpec.implementation)) {
+    return componentSpec;
+  }
+
+  const graphSpec = componentSpec.implementation.graph;
+  const updatedComponentSpec = { ...componentSpec };
 
-  if (!edge.targetHandle) {
+  const targetNodeId = edge.target;
+  const targetTaskId = nodeManager.getTaskId(targetNodeId);
+  const targetNodeType = nodeManager.getNodeType(targetNodeId);
+
+  if (!targetTaskId || !targetNodeType) {
+    console.error("Could not resolve target node information:", {
+      targetNodeId,
+      targetTaskId,
+      targetNodeType,
+    });
     return componentSpec;
   }
 
-  const inputName = nodeManager.getHandleInfo(edge.targetHandle)?.handleName;
+  switch (targetNodeType) {
+    case "task": {
+      if (!edge.targetHandle) {
+        console.error("No target handle found for task connection");
+        return componentSpec;
+      }
+
+      const targetHandleInfo = nodeManager.getHandleInfo(edge.targetHandle);
+      if (!targetHandleInfo) {
+        console.error("Could not resolve target handle info");
+        return componentSpec;
+      }
 
-  const updatedComponentSpec = {
-    ...componentSpec,
-  };
+      const inputName = targetHandleInfo.handleName;
+      const newGraphSpec = setTaskArgument(graphSpec, targetTaskId, inputName);
 
-  const taskId = nodeManager.getTaskId(edge.target);
-  if (!taskId) return componentSpec;
+      updatedComponentSpec.implementation = {
+        ...updatedComponentSpec.implementation,
+        graph: newGraphSpec,
+      };
+      break;
+    }
 
-  if (inputName !== undefined && graphSpec) {
-    const newGraphSpec = setTaskArgument(graphSpec, taskId, inputName);
-    updatedComponentSpec.implementation = {
-      ...updatedComponentSpec.implementation,
-      graph: newGraphSpec,
-    };
+    case "output": {
+      const outputName = outputIdToOutputName(targetTaskId);
+      const newGraphSpec = setGraphOutputValue(graphSpec, outputName);
 
-    return updatedComponentSpec;
-  } else {
-    const outputName = outputIdToOutputName(taskId);
-    const newGraphSpec = setGraphOutputValue(graphSpec, outputName);
-    updatedComponentSpec.implementation = {
-      ...updatedComponentSpec.implementation,
-      graph: newGraphSpec,
-    };
+      updatedComponentSpec.implementation = {
+        ...updatedComponentSpec.implementation,
+        graph: newGraphSpec,
+      };
+      break;
+    }
 
-    return updatedComponentSpec;
+    default:
+      console.error(
+        "Unsupported target node type for edge removal:",
+        targetNodeType,
+      );
+      return componentSpec;
   }
+
+  return updatedComponentSpec;
 };
diff --git a/src/hooks/useComponentSpecToEdges.ts b/src/hooks/useComponentSpecToEdges.ts
