fix+refactor(builder): Sort out FlowEditor+ReactFlow state management (#7737)

* refactor(builder): Migrate `FlowEditor` to use ReactFlow's state management system We have been keeping two copies of node and edge data: one inside ReactFlow and one outside. It works, but it's accidental and implicit and there is no reason to be using shadow copies rather than a single data source. - Replace `useNodesState` and `useEdgesState` with `useReactFlow` hook - Use `addNodes`, `addEdges`, and `deleteElements` where appropriate instead of `setNodes`/`setEdges` to allow use of event hooks - Consolidate all edge -> node state sync logic into `onEdgesChange` event handler This replaces `updateNodesOnEdgeChange`, part of `onConnect`, and `onEdgesDelete`. - Move node deletion logic from `CustomNode` to `FlowEditor:onNodesChange` * fix(builder): Refactor and fix copy-paste mechanism - Rename variables for readability - Use an ID map to correctly set the source and target IDs for the pasted edges
2026-04-08 03:00:28 -04:00 · 2024-08-08 16:18:15 +02:00
parent bf10df612e
commit 582571631e
4 changed files with 240 additions and 239 deletions
--- a/rnd/autogpt_builder/src/components/CustomEdge.tsx
+++ b/rnd/autogpt_builder/src/components/CustomEdge.tsx
@@ -30,10 +30,10 @@ const CustomEdgeFC: FC<EdgeProps<CustomEdgeData>> = ({
  markerEnd,
 }) => {
  const [isHovered, setIsHovered] = useState(false);
-  const { setEdges } = useReactFlow();
+  const { deleteElements } = useReactFlow<any, CustomEdgeData>();

  const onEdgeRemoveClick = () => {
-    setEdges((edges) => edges.filter((edge) => edge.id !== id));
+    deleteElements({ edges: [{ id }] });
  };

  const [path, labelX, labelY] = getBezierPath({
--- a/rnd/autogpt_builder/src/components/CustomNode.tsx
+++ b/rnd/autogpt_builder/src/components/CustomNode.tsx
@@ -34,6 +34,7 @@ export type CustomNodeData = {
  hardcodedValues: { [key: string]: any };
  setHardcodedValues: (values: { [key: string]: any }) => void;
  connections: Array<{
+    edge_id: string;
    source: string;
    sourceHandle: string;
    target: string;
@@ -58,10 +59,7 @@ const CustomNode: FC<NodeProps<CustomNodeData>> = ({ data, id }) => {
  const [isOutputModalOpen, setIsOutputModalOpen] = useState(false);
  const [isHovered, setIsHovered] = useState(false);

-  const { getNode, setNodes, getEdges, setEdges } = useReactFlow<
-    CustomNodeData,
-    CustomEdgeData
-  >();
+  const { deleteElements } = useReactFlow();

  const outputDataRef = useRef<HTMLDivElement>(null);
  const isInitialSetup = useRef(true);
@@ -267,44 +265,9 @@ const CustomNode: FC<NodeProps<CustomNodeData>> = ({ data, id }) => {
  const deleteNode = useCallback(() => {
    console.log("Deleting node:", id);

-    // Get all edges connected to this node
-    const connectedEdges = getEdges().filter(
-      (edge) => edge.source === id || edge.target === id,
-    );
-
-    // For each connected edge, update the connected node's state
-    connectedEdges.forEach((edge) => {
-      const connectedNodeId = edge.source === id ? edge.target : edge.source;
-      const connectedNode = getNode(connectedNodeId);
-
-      if (connectedNode) {
-        setNodes((nodes) =>
-          nodes.map((node) => {
-            if (node.id === connectedNodeId) {
-              // Update the node's data to reflect the disconnection
-              const updatedConnections = node.data.connections.filter(
-                (conn) => !(conn.source === id || conn.target === id),
-              );
-              return {
-                ...node,
-                data: {
-                  ...node.data,
-                  connections: updatedConnections,
-                },
-              };
-            }
-            return node;
-          }),
-        );
-      }
-    });
-
-    // Remove the node and its connected edges
-    setNodes((nodes) => nodes.filter((node) => node.id !== id));
-    setEdges((edges) =>
-      edges.filter((edge) => edge.source !== id && edge.target !== id),
-    );
-  }, [id, setNodes, setEdges, getNode, getEdges]);
+    // Remove the node
+    deleteElements({ nodes: [{ id }] });
+  }, [id, deleteElements]);

  const copyNode = useCallback(() => {
    // This is a placeholder function. The actual copy functionality
--- a/rnd/autogpt_builder/src/components/Flow.tsx
+++ b/rnd/autogpt_builder/src/components/Flow.tsx
@@ -5,11 +5,13 @@ import React, {
  useEffect,
  useMemo,
  useRef,
+  MouseEvent,
 } from "react";
+import { shallow } from "zustand/vanilla/shallow";
 import ReactFlow, {
-  addEdge,
-  useNodesState,
-  useEdgesState,
+  ReactFlowProvider,
+  Controls,
+  Background,
  Node,
  Edge,
  OnConnect,
@@ -17,15 +19,21 @@ import ReactFlow, {
  Connection,
  EdgeTypes,
  MarkerType,
-  Controls,
-  Background,
+  NodeChange,
+  EdgeChange,
+  useStore,
+  useReactFlow,
+  applyEdgeChanges,
+  applyNodeChanges,
 } from "reactflow";
 import "reactflow/dist/style.css";
 import CustomNode, { CustomNodeData } from "./CustomNode";
 import "./flow.css";
 import AutoGPTServerAPI, {
  Block,
+  BlockIOSubSchema,
  Graph,
+  Link,
  NodeExecutionResult,
 } from "@/lib/autogpt-server-api";
 import {
@@ -54,8 +62,26 @@ const FlowEditor: React.FC<{
  template?: boolean;
  className?: string;
 }> = ({ flowID, template, className }) => {
-  const [nodes, setNodes, onNodesChange] = useNodesState<CustomNodeData>([]);
-  const [edges, setEdges, onEdgesChange] = useEdgesState<CustomEdgeData>([]);
+  const { _setNodes, _setEdges } = useStore(
+    useCallback(
+      ({ setNodes, setEdges }) => ({
+        _setNodes: setNodes,
+        _setEdges: setEdges,
+      }),
+      [],
+    ),
+    shallow,
+  );
+  const {
+    addNodes,
+    addEdges,
+    getNode,
+    getNodes,
+    getEdges,
+    setNodes,
+    setEdges,
+    deleteElements,
+  } = useReactFlow<CustomNodeData, CustomEdgeData>();
  const [nodeId, setNodeId] = useState<number>(1);
  const [availableNodes, setAvailableNodes] = useState<Block[]>([]);
  const [savedAgent, setSavedAgent] = useState<Graph | null>(null);
@@ -136,12 +162,12 @@ const FlowEditor: React.FC<{
  const nodeTypes: NodeTypes = useMemo(() => ({ custom: CustomNode }), []);
  const edgeTypes: EdgeTypes = useMemo(() => ({ custom: CustomEdge }), []);

-  const onNodesChangeStart = (event: MouseEvent, node: Node) => {
+  const onNodeDragStart = (_: MouseEvent, node: Node) => {
    initialPositionRef.current[node.id] = { ...node.position };
    isDragging.current = true;
  };

-  const onNodesChangeEnd = (event: MouseEvent, node: Node | null) => {
+  const onNodeDragEnd = (_: MouseEvent, node: Node | null) => {
    if (!node) return;

    isDragging.current = false;
@@ -162,12 +188,14 @@ const FlowEditor: React.FC<{
        type: "UPDATE_NODE_POSITION",
        payload: { nodeId: node.id, oldPosition, newPosition },
        undo: () =>
+          // TODO: replace with updateNodes() after upgrade to ReactFlow v12
          setNodes((nds) =>
            nds.map((n) =>
              n.id === node.id ? { ...n, position: oldPosition } : n,
            ),
          ),
        redo: () =>
+          // TODO: replace with updateNodes() after upgrade to ReactFlow v12
          setNodes((nds) =>
            nds.map((n) =>
              n.id === node.id ? { ...n, position: newPosition } : n,
@@ -178,48 +206,8 @@ const FlowEditor: React.FC<{
    delete initialPositionRef.current[node.id];
  };

-  const updateNodesOnEdgeChange = (
-    edge: Edge<CustomEdgeData>,
-    action: "add" | "remove",
-  ) => {
-    setNodes((nds) =>
-      nds.map((node) => {
-        if (node.id === edge.source || node.id === edge.target) {
-          const connections =
-            action === "add"
-              ? [
-                  ...node.data.connections,
-                  {
-                    source: edge.source,
-                    sourceHandle: edge.sourceHandle!,
-                    target: edge.target,
-                    targetHandle: edge.targetHandle!,
-                  },
-                ]
-              : node.data.connections.filter(
-                  (conn) =>
-                    !(
-                      conn.source === edge.source &&
-                      conn.target === edge.target &&
-                      conn.sourceHandle === edge.sourceHandle &&
-                      conn.targetHandle === edge.targetHandle
-                    ),
-                );
-          return {
-            ...node,
-            data: {
-              ...node.data,
-              connections,
-            },
-          };
-        }
-        return node;
-      }),
-    );
-  };
-
  const getOutputType = (id: string, handleId: string) => {
-    const node = nodes.find((node) => node.id === id);
+    const node = getNode(id);
    if (!node) return "unknown";

    const outputSchema = node.data.outputSchema;
@@ -230,13 +218,6 @@ const FlowEditor: React.FC<{
    return outputHandle.type;
  };

-  const getNodePos = (id: string) => {
-    const node = nodes.find((node) => node.id === id);
-    if (!node) return 0;
-
-    return node.position;
-  };
-
  // Function to clear status, output, and close the output info dropdown of all nodes
  const clearNodesStatusAndOutput = useCallback(() => {
    setNodes((nds) =>
@@ -252,15 +233,37 @@ const FlowEditor: React.FC<{
    );
  }, [setNodes]);

+  const onNodesChange = useCallback(
+    (nodeChanges: NodeChange[]) => {
+      // Persist the changes
+      _setNodes(applyNodeChanges(nodeChanges, getNodes()));
+
+      // Remove all edges that were connected to deleted nodes
+      nodeChanges
+        .filter((change) => change.type == "remove")
+        .forEach((deletedNode) => {
+          const nodeID = deletedNode.id;
+
+          const connectedEdges = getEdges().filter((edge) =>
+            [edge.source, edge.target].includes(nodeID),
+          );
+          deleteElements({
+            edges: connectedEdges.map((edge) => ({ id: edge.id })),
+          });
+        });
+    },
+    [getNodes, getEdges, _setNodes, deleteElements],
+  );
+
  const onConnect: OnConnect = useCallback(
    (connection: Connection) => {
      const edgeColor = getTypeColor(
        getOutputType(connection.source!, connection.sourceHandle!),
      );
-      const sourcePos = getNodePos(connection.source!);
+      const sourcePos = getNode(connection.source!)?.position;
      console.log("sourcePos", sourcePos);
-      const newEdge = {
-        id: `${connection.source}_${connection.sourceHandle}_${connection.target}_${connection.targetHandle}`,
+      const newEdge: Edge<CustomEdgeData> = {
+        id: formatEdgeID(connection),
        type: "custom",
        markerEnd: {
          type: MarkerType.ArrowClosed,
@@ -269,82 +272,97 @@ const FlowEditor: React.FC<{
        },
        data: { edgeColor, sourcePos },
        ...connection,
+        source: connection.source!,
+        target: connection.target!,
      };

-      setEdges((eds) => {
-        const newEdges = addEdge(newEdge, eds);
-        history.push({
-          type: "ADD_EDGE",
-          payload: newEdge,
-          undo: () => {
-            setEdges((prevEdges) =>
-              prevEdges.filter((edge) => edge.id !== newEdge.id),
-            );
-            updateNodesOnEdgeChange(newEdge, "remove");
-          },
-          redo: () => {
-            setEdges((prevEdges) => addEdge(newEdge, prevEdges));
-            updateNodesOnEdgeChange(newEdge, "add");
-          },
-        });
-        updateNodesOnEdgeChange(newEdge, "add");
-        return newEdges;
+      addEdges(newEdge);
+      history.push({
+        type: "ADD_EDGE",
+        payload: { edge: newEdge },
+        undo: () => {
+          deleteElements({ edges: [{ id: newEdge.id }] });
+        },
+        redo: () => {
+          addEdges(newEdge);
+        },
      });
-      setNodes((nds) =>
-        nds.map((node) => {
-          if (node.id === connection.target || node.id === connection.source) {
-            return {
-              ...node,
-              data: {
-                ...node.data,
-                connections: [
-                  ...node.data.connections,
-                  {
-                    source: connection.source,
-                    sourceHandle: connection.sourceHandle,
-                    target: connection.target,
-                    targetHandle: connection.targetHandle,
-                  } as {
-                    source: string;
-                    sourceHandle: string;
-                    target: string;
-                    targetHandle: string;
-                  },
-                ],
-              },
-            };
-          }
-          return node;
-        }),
-      );
      clearNodesStatusAndOutput(); // Clear status and output on connection change
    },
-    [nodes],
+    [getNode, addEdges, history, deleteElements, clearNodesStatusAndOutput],
  );

-  const onEdgesDelete = useCallback(
-    (edgesToDelete: Edge<CustomEdgeData>[]) => {
-      setNodes((nds) =>
-        nds.map((node) => ({
-          ...node,
-          data: {
-            ...node.data,
-            connections: node.data.connections.filter(
-              (conn: any) =>
-                !edgesToDelete.some(
-                  (edge) =>
-                    edge.source === conn.source &&
-                    edge.target === conn.target &&
-                    edge.sourceHandle === edge.sourceHandle &&
-                    edge.targetHandle === edge.targetHandle,
+  const onEdgesChange = useCallback(
+    (edgeChanges: EdgeChange[]) => {
+      // Persist the changes
+      _setEdges(applyEdgeChanges(edgeChanges, getEdges()));
+
+      // Propagate edge changes to node data
+      const addedEdges = edgeChanges.filter((change) => change.type == "add"),
+        resetEdges = edgeChanges.filter((change) => change.type == "reset"),
+        removedEdges = edgeChanges.filter((change) => change.type == "remove"),
+        selectedEdges = edgeChanges.filter((change) => change.type == "select");
+
+      if (addedEdges.length > 0 || removedEdges.length > 0) {
+        setNodes((nds) =>
+          nds.map((node) => ({
+            ...node,
+            data: {
+              ...node.data,
+              connections: [
+                // Remove node connections for deleted edges
+                ...node.data.connections.filter(
+                  (conn) =>
+                    !removedEdges.some(
+                      (removedEdge) => removedEdge.id == conn.edge_id,
+                    ),
                ),
-            ),
-          },
-        })),
-      );
-      clearNodesStatusAndOutput(); // Clear status and output on edge deletion
+                // Add node connections for added edges
+                ...addedEdges.map((addedEdge) => ({
+                  edge_id: addedEdge.item.id,
+                  source: addedEdge.item.source,
+                  target: addedEdge.item.target,
+                  sourceHandle: addedEdge.item.sourceHandle!,
+                  targetHandle: addedEdge.item.targetHandle!,
+                })),
+              ],
+            },
+          })),
+        );
+
+        if (removedEdges.length > 0) {
+          clearNodesStatusAndOutput(); // Clear status and output on edge deletion
+        }
+      }
+
+      if (resetEdges.length > 0) {
+        // Reset node connections for all edges
+        console.warn(
+          "useReactFlow().setEdges was used to overwrite all edges. " +
+            "Use addEdges, deleteElements, or reconnectEdge for incremental changes.",
+          resetEdges,
+        );
+        setNodes((nds) =>
+          nds.map((node) => ({
+            ...node,
+            data: {
+              ...node.data,
+              connections: [
+                ...resetEdges.map((resetEdge) => ({
+                  edge_id: resetEdge.item.id,
+                  source: resetEdge.item.source,
+                  target: resetEdge.item.target,
+                  sourceHandle: resetEdge.item.sourceHandle!,
+                  targetHandle: resetEdge.item.targetHandle!,
+                })),
+              ],
+            },
+          })),
+        );
+        clearNodesStatusAndOutput();
+      }
    },
-    [setNodes, clearNodesStatusAndOutput],
+    [getEdges, _setEdges, setNodes, clearNodesStatusAndOutput],
  );

  const addNode = useCallback(
@@ -366,6 +384,7 @@ const FlowEditor: React.FC<{
          outputSchema: nodeSchema.outputSchema,
          hardcodedValues: {},
          setHardcodedValues: (values) => {
+            // TODO: replace with updateNodes() after upgrade to ReactFlow v12
            setNodes((nds) =>
              nds.map((node) =>
                node.id === newNode.id
@@ -377,8 +396,9 @@ const FlowEditor: React.FC<{
          connections: [],
          isOutputOpen: false,
          block_id: blockId,
-          setIsAnyModalOpen: setIsAnyModalOpen, // Pass setIsAnyModalOpen function
+          setIsAnyModalOpen,
          setErrors: (errors: { [key: string]: string | null }) => {
+            // TODO: replace with updateNodes() after upgrade to ReactFlow v12
            setNodes((nds) =>
              nds.map((node) =>
                node.id === newNode.id
@@ -390,19 +410,25 @@ const FlowEditor: React.FC<{
        },
      };

-      setNodes((nds) => [...nds, newNode]);
+      addNodes(newNode);
      setNodeId((prevId) => prevId + 1);
      clearNodesStatusAndOutput(); // Clear status and output when a new node is added

      history.push({
        type: "ADD_NODE",
-        payload: newNode,
-        undo: () =>
-          setNodes((nds) => nds.filter((node) => node.id !== newNode.id)),
-        redo: () => setNodes((nds) => [...nds, newNode]),
+        payload: { node: newNode.data },
+        undo: () => deleteElements({ nodes: [{ id: newNode.id }] }),
+        redo: () => addNodes(newNode),
      });
    },
-    [nodeId, availableNodes],
+    [
+      nodeId,
+      availableNodes,
+      addNodes,
+      setNodes,
+      deleteElements,
+      clearNodesStatusAndOutput,
+    ],
  );

  const handleUndo = () => {
@@ -431,7 +457,6 @@ const FlowEditor: React.FC<{
            y: node.metadata.position.y,
          },
          data: {
-            setIsAnyModalOpen: setIsAnyModalOpen,
            block_id: block.id,
            blockType: block.name,
            title: `${block.name} ${node.id}`,
@@ -453,12 +478,14 @@ const FlowEditor: React.FC<{
            connections: graph.links
              .filter((l) => [l.source_id, l.sink_id].includes(node.id))
              .map((link) => ({
+                edge_id: formatEdgeID(link),
                source: link.source_id,
                sourceHandle: link.source_name,
                target: link.sink_id,
                targetHandle: link.sink_name,
              })),
            isOutputOpen: false,
+            setIsAnyModalOpen,
            setErrors: (errors: { [key: string]: string | null }) => {
              setNodes((nds) =>
                nds.map((node) =>
@@ -475,30 +502,25 @@ const FlowEditor: React.FC<{
    );

    setEdges(
-      graph.links.map(
-        (link) =>
-          ({
-            id: `${link.source_id}_${link.source_name}_${link.sink_id}_${link.sink_name}`,
-            type: "custom",
-            data: {
-              edgeColor: getTypeColor(
-                getOutputType(link.source_id, link.source_name!),
-              ),
-              sourcePos: getNodePos(link.source_id),
-            },
-            markerEnd: {
-              type: MarkerType.ArrowClosed,
-              strokeWidth: 2,
-              color: getTypeColor(
-                getOutputType(link.source_id, link.source_name!),
-              ),
-            },
-            source: link.source_id,
-            target: link.sink_id,
-            sourceHandle: link.source_name || undefined,
-            targetHandle: link.sink_name || undefined,
-          }) as Edge<CustomEdgeData>,
-      ),
+      graph.links.map((link) => ({
+        id: formatEdgeID(link),
+        type: "custom",
+        data: {
+          edgeColor: getTypeColor(
+            getOutputType(link.source_id, link.source_name),
+          ),
+          sourcePos: getNode(link.source_id)?.position,
+        },
+        markerEnd: {
+          type: MarkerType.ArrowClosed,
+          strokeWidth: 2,
+          color: getTypeColor(getOutputType(link.source_id, link.source_name)),
+        },
+        source: link.source_id,
+        target: link.sink_id,
+        sourceHandle: link.source_name || undefined,
+        targetHandle: link.sink_name || undefined,
+      })),
    );
  }

@@ -568,8 +590,11 @@ const FlowEditor: React.FC<{
      })),
    );
    await new Promise((resolve) => setTimeout(resolve, 100));
+
+    const nodes = getNodes();
+    const edges = getEdges();
    console.log("All nodes before formatting:", nodes);
-    const blockIdToNodeIdMap = {};
+    const blockIdToNodeIdMap: Record<string, string> = {};

    const formattedNodes = nodes.map((node) => {
      nodes.forEach((node) => {
@@ -671,7 +696,7 @@ const FlowEditor: React.FC<{
  const validateNodes = (): boolean => {
    let isValid = true;

-    nodes.forEach((node) => {
+    getNodes().forEach((node) => {
      const validate = ajv.compile(node.data.inputSchema);
      const errors = {} as { [key: string]: string | null };

@@ -765,16 +790,18 @@ const FlowEditor: React.FC<{
      if (event.ctrlKey || event.metaKey) {
        if (event.key === "c" || event.key === "C") {
          // Copy selected nodes
-          const selectedNodes = nodes.filter((node) => node.selected);
-          const selectedEdges = edges.filter((edge) => edge.selected);
+          const selectedNodes = getNodes().filter((node) => node.selected);
+          const selectedEdges = getEdges().filter((edge) => edge.selected);
          setCopiedNodes(selectedNodes);
          setCopiedEdges(selectedEdges);
        }
        if (event.key === "v" || event.key === "V") {
          // Paste copied nodes
          if (copiedNodes.length > 0) {
-            const newNodes = copiedNodes.map((node, index) => {
+            const oldToNewNodeIDMap: Record<string, string> = {};
+            const pastedNodes = copiedNodes.map((node, index) => {
              const newNodeId = (nodeId + index).toString();
+              oldToNewNodeIDMap[node.id] = newNodeId;
              return {
                ...node,
                id: newNodeId,
@@ -801,20 +828,16 @@ const FlowEditor: React.FC<{
                },
              };
            });
-            const updatedNodes = nodes.map((node) => ({
-              ...node,
-              selected: false,
-            })); // Deselect old nodes
-            setNodes([...updatedNodes, ...newNodes]);
+            setNodes((existingNodes) =>
+              // Deselect copied nodes
+              existingNodes.map((node) => ({ ...node, selected: false })),
+            );
+            addNodes(pastedNodes);
            setNodeId((prevId) => prevId + copiedNodes.length);

-            const newEdges = copiedEdges.map((edge) => {
-              const newSourceId =
-                newNodes.find((n) => n.data.title === edge.source)?.id ||
-                edge.source;
-              const newTargetId =
-                newNodes.find((n) => n.data.title === edge.target)?.id ||
-                edge.target;
+            const pastedEdges = copiedEdges.map((edge) => {
+              const newSourceId = oldToNewNodeIDMap[edge.source] ?? edge.source;
+              const newTargetId = oldToNewNodeIDMap[edge.target] ?? edge.target;
              return {
                ...edge,
                id: `${newSourceId}_${edge.sourceHandle}_${newTargetId}_${edge.targetHandle}_${Date.now()}`,
@@ -822,12 +845,22 @@ const FlowEditor: React.FC<{
                target: newTargetId,
              };
            });
-            setEdges([...edges, ...newEdges]);
+            addEdges(pastedEdges);
          }
        }
      }
    },
-    [nodes, edges, copiedNodes, copiedEdges, nodeId, isAnyModalOpen],
+    [
+      addNodes,
+      addEdges,
+      getNodes,
+      getEdges,
+      setNodes,
+      copiedNodes,
+      copiedEdges,
+      nodeId,
+      isAnyModalOpen,
+    ],
  );

  useEffect(() => {
@@ -862,25 +895,16 @@ const FlowEditor: React.FC<{
  return (
    <div className={className}>
      <ReactFlow
-        nodes={nodes.map((node) => ({
-          ...node,
-          data: { ...node.data, setIsAnyModalOpen },
-        }))}
-        edges={edges.map((edge) => ({
-          ...edge,
-          data: { ...edge.data, clearNodesStatusAndOutput },
-        }))}
-        onNodesChange={onNodesChange}
-        onEdgesChange={onEdgesChange}
-        onConnect={onConnect}
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        connectionLineComponent={ConnectionLine}
+        onConnect={onConnect}
+        onNodesChange={onNodesChange}
        onNodesDelete={onNodesDelete}
-        onEdgesDelete={onEdgesDelete}
+        onEdgesChange={onEdgesChange}
+        onNodeDragStop={onNodeDragEnd}
+        onNodeDragStart={onNodeDragStart}
        deleteKeyCode={["Backspace", "Delete"]}
-        onNodeDragStart={onNodesChangeStart}
-        onNodeDragStop={onNodesChangeEnd}
      >
        <Controls />
        <Background />
@@ -898,4 +922,18 @@ const FlowEditor: React.FC<{
  );
 };

-export default FlowEditor;
+const WrappedFlowEditor: typeof FlowEditor = (props) => (
+  <ReactFlowProvider>
+    <FlowEditor {...props} />
+  </ReactFlowProvider>
+);
+
+export default WrappedFlowEditor;
+
+function formatEdgeID(conn: Link | Connection): string {
+  if ("sink_id" in conn) {
+    return `${conn.source_id}_${conn.source_name}_${conn.sink_id}_${conn.sink_name}`;
+  } else {
+    return `${conn.source}_${conn.sourceHandle}_${conn.target}_${conn.targetHandle}`;
+  }
+}
--- a/rnd/autogpt_builder/src/lib/autogpt-server-api/types.ts
+++ b/rnd/autogpt_builder/src/lib/autogpt-server-api/types.ts
@@ -10,7 +10,7 @@ export type Block = {
 export type BlockIORootSchema = {
  type: "object";
  properties: { [key: string]: BlockIOSubSchema };
-  required?: string[];
+  required?: (keyof BlockIORootSchema["properties"])[];
  additionalProperties?: { type: string };
 };

@@ -37,7 +37,7 @@ export type BlockIOObjectSubSchema = BlockIOSubSchemaMeta & {
  type: "object";
  properties: { [key: string]: BlockIOSubSchema };
  default?: { [key: keyof BlockIOObjectSubSchema["properties"]]: any };
-  required?: keyof BlockIOObjectSubSchema["properties"][];
+  required?: (keyof BlockIOObjectSubSchema["properties"])[];
 };

 export type BlockIOKVSubSchema = BlockIOSubSchemaMeta & {