From d9a1a1edc8d590490c614041dedcda0946d470bb Mon Sep 17 00:00:00 2001
From: Reinier van der Leer <pwuts@agpt.co>
Date: Thu, 25 Jul 2024 10:30:35 +0200
Subject: [PATCH] fix(builder): Correctly load node connections (#7592)

- Set `node.data.connections` based on `graph.links` in `loadGraph(..)`
- Use reactflow's built-in `useNodesState`, `useEdgesState` to replace some boilerplate functions that do exactly the same

Resolves #7392
---
 .../src/components/CustomNode.tsx             |  2 +-
 rnd/autogpt_builder/src/components/Flow.tsx   | 31 ++++++++-----------
 2 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/rnd/autogpt_builder/src/components/CustomNode.tsx b/rnd/autogpt_builder/src/components/CustomNode.tsx
index e3c4fff6f..5fbeea15d 100644
--- a/rnd/autogpt_builder/src/components/CustomNode.tsx
+++ b/rnd/autogpt_builder/src/components/CustomNode.tsx
@@ -167,7 +167,7 @@ const CustomNode: FC<NodeProps<CustomNodeData>> = ({ data, id }) => {
               return (isRequired || isAdvancedOpen) && (
                 <div key={key}>
                   <NodeHandle keyName={key} isConnected={isHandleConnected(key)} schema={schema} side="left" />
-                  {isHandleConnected(key) ? <></> :
+                  {!isHandleConnected(key) &&
                   <NodeInputField
                     keyName={key}
                     schema={schema}
diff --git a/rnd/autogpt_builder/src/components/Flow.tsx b/rnd/autogpt_builder/src/components/Flow.tsx
index 739d954e3..730d14298 100644
--- a/rnd/autogpt_builder/src/components/Flow.tsx
+++ b/rnd/autogpt_builder/src/components/Flow.tsx
@@ -2,12 +2,10 @@
 import React, { useState, useCallback, useEffect, useMemo } from 'react';
 import ReactFlow, {
   addEdge,
-  applyNodeChanges,
-  applyEdgeChanges,
+  useNodesState,
+  useEdgesState,
   Node,
   Edge,
-  OnNodesChange,
-  OnEdgesChange,
   OnConnect,
   NodeTypes,
   Connection,
@@ -76,8 +74,8 @@ const FlowEditor: React.FC<{
   template?: boolean;
   className?: string;
 }> = ({ flowID, template, className }) => {
-  const [nodes, setNodes] = useState<Node<CustomNodeData>[]>([]);
-  const [edges, setEdges] = useState<Edge<CustomEdgeData>[]>([]);
+  const [nodes, setNodes, onNodesChange] = useNodesState<CustomNodeData>([]);
+  const [edges, setEdges, onEdgesChange] = useEdgesState<CustomEdgeData>([]);
   const [nodeId, setNodeId] = useState<number>(1);
   const [availableNodes, setAvailableNodes] = useState<Block[]>([]);
   const [isSidebarOpen, setIsSidebarOpen] = useState(true);
@@ -124,16 +122,6 @@ const FlowEditor: React.FC<{
   const nodeTypes: NodeTypes = useMemo(() => ({ custom: CustomNode }), []);
   const edgeTypes: EdgeTypes = useMemo(() => ({ custom: CustomEdge }), []);
 
-  const onNodesChange: OnNodesChange = useCallback(
-    (changes) => setNodes((nds) => applyNodeChanges(changes, nds)),
-    []
-  );
-
-  const onEdgesChange: OnEdgesChange = useCallback(
-    (changes) => setEdges((eds) => applyEdgeChanges(changes, eds)),
-    []
-  );
-
   const getOutputType = (id: string, handleId: string) => {
     const node = nodes.find((node) => node.id === id);
     if (!node) return 'unknown';
@@ -251,7 +239,7 @@ const FlowEditor: React.FC<{
 
     setNodes(graph.nodes.map(node => {
       const block = availableNodes.find(block => block.id === node.block_id)!;
-      const newNode = {
+      const newNode: Node<CustomNodeData> = {
         id: node.id,
         type: 'custom',
         position: { x: node.metadata.position.x, y: node.metadata.position.y },
@@ -268,7 +256,14 @@ const FlowEditor: React.FC<{
               : node
             ));
           },
-          connections: [],
+          connections: graph.links
+            .filter(l => [l.source_id, l.sink_id].includes(node.id))
+            .map(link => ({
+              source: link.source_id,
+              sourceHandle: link.source_name,
+              target: link.sink_id,
+              targetHandle: link.sink_name,
+            })),
           isOutputOpen: false,
         },
       };