diff --git a/src/components/shared/ReactFlow/FlowCanvas/utils/removeEdge.ts b/src/components/shared/ReactFlow/FlowCanvas/utils/removeEdge.ts
index 3e02034a7..4a6f3fecd 100644
--- 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,34 +15,67 @@ export const removeEdge = (
   componentSpec: ComponentSpec,
   nodeManager: NodeManager,
 ) => {
-  const graphSpec = (componentSpec.implementation as GraphImplementation)
-    ?.graph;
-
-  const inputName = edge.targetHandle?.replace(/^input_/, "");
-
-  const updatedComponentSpec = {
-    ...componentSpec,
-  };
-
-  const taskId = nodeManager.getTaskId(edge.target);
-  if (!taskId) return componentSpec;
-
-  if (inputName !== undefined && graphSpec) {
-    const newGraphSpec = setTaskArgument(graphSpec, taskId, inputName);
-    updatedComponentSpec.implementation = {
-      ...updatedComponentSpec.implementation,
-      graph: newGraphSpec,
-    };
-
-    return updatedComponentSpec;
-  } else {
-    const outputName = outputIdToOutputName(taskId);
-    const newGraphSpec = setGraphOutputValue(graphSpec, outputName);
-    updatedComponentSpec.implementation = {
-      ...updatedComponentSpec.implementation,
-      graph: newGraphSpec,
-    };
-
-    return updatedComponentSpec;
+  if (!isGraphImplementation(componentSpec.implementation)) {
+    return componentSpec;
   }
+
+  const graphSpec = componentSpec.implementation.graph;
+  const updatedComponentSpec = { ...componentSpec };
+
+  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;
+  }
+
+  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 inputName = targetHandleInfo.handleName;
+      const newGraphSpec = setTaskArgument(graphSpec, targetTaskId, inputName);
+
+      updatedComponentSpec.implementation = {
+        ...updatedComponentSpec.implementation,
+        graph: newGraphSpec,
+      };
+      break;
+    }
+
+    case "output": {
+      const outputName = outputIdToOutputName(targetTaskId);
+      const newGraphSpec = setGraphOutputValue(graphSpec, outputName);
+
+      updatedComponentSpec.implementation = {
+        ...updatedComponentSpec.implementation,
+        graph: newGraphSpec,
+      };
+      break;
+    }
+
+    default:
+      console.error(
+        "Unsupported target node type for edge removal:",
+        targetNodeType,
+      );
+      return componentSpec;
+  }
+
+  return updatedComponentSpec;
 };