diff --git a/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/Flow.tsx b/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/Flow.tsx index 29e120f348..8d4963dd63 100644 --- a/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/Flow.tsx +++ b/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/Flow.tsx @@ -20,6 +20,7 @@ export const Flow = () => { useShallow((state) => state.onNodesChange), ); const nodeTypes = useMemo(() => ({ custom: CustomNode }), []); + const edgeTypes = useMemo(() => ({ custom: CustomEdge }), []); const { edges, onConnect, onEdgesChange } = useCustomEdge(); // We use this hook to load the graph and convert them into custom nodes and edges. @@ -51,10 +52,10 @@ export const Flow = () => { nodes={nodes} onNodesChange={onNodesChange} nodeTypes={nodeTypes} + edgeTypes={edgeTypes} edges={edges} onConnect={onConnect} onEdgesChange={onEdgesChange} - edgeTypes={{ custom: CustomEdge }} maxZoom={2} minZoom={0.1} onDragOver={onDragOver} diff --git a/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/useFlow.ts b/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/useFlow.ts index 069daa3aea..6d65e0bd01 100644 --- a/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/useFlow.ts +++ b/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/useFlow.ts @@ -32,6 +32,9 @@ export const useFlow = () => { const setGraphSchemas = useGraphStore( useShallow((state) => state.setGraphSchemas), ); + const updateEdgeBeads = useEdgeStore( + useShallow((state) => state.updateEdgeBeads), + ); const { screenToFlowPosition } = useReactFlow(); const addBlock = useNodeStore(useShallow((state) => state.addBlock)); const setBlockMenuOpen = useControlPanelStore( @@ -109,7 +112,7 @@ export const useFlow = () => { // adding links if (graph?.links) { - useEdgeStore.getState().setConnections([]); + useEdgeStore.getState().setEdges([]); addLinks(graph.links); } @@ -130,7 +133,7 @@ export const useFlow = () => { }); } - // update node execution results in nodes + // update node execution results in nodes, also update edge beads if ( executionDetails && "node_executions" in executionDetails && @@ -138,6 +141,7 @@ export const useFlow = () => { ) { executionDetails.node_executions.forEach((nodeExecution) => { updateNodeExecutionResult(nodeExecution.node_id, nodeExecution); + updateEdgeBeads(nodeExecution.node_id, nodeExecution); }); } }, [customNodes, addNodes, graph?.links, executionDetails, updateNodeStatus]); @@ -145,8 +149,9 @@ export const useFlow = () => { useEffect(() => { return () => { useNodeStore.getState().setNodes([]); - useEdgeStore.getState().setConnections([]); + useEdgeStore.getState().setEdges([]); useGraphStore.getState().reset(); + useEdgeStore.getState().resetEdgeBeads(); setIsGraphRunning(false); }; }, []); diff --git a/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/useFlowRealtime.ts b/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/useFlowRealtime.ts index 7ab55554f5..081a7048c5 100644 --- a/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/useFlowRealtime.ts +++ b/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/Flow/useFlowRealtime.ts @@ -9,6 +9,7 @@ import { useShallow } from "zustand/react/shallow"; import { NodeExecutionResult } from "@/app/api/__generated__/models/nodeExecutionResult"; import { AgentExecutionStatus } from "@/app/api/__generated__/models/agentExecutionStatus"; import { useGraphStore } from "../../../stores/graphStore"; +import { useEdgeStore } from "../../../stores/edgeStore"; export const useFlowRealtime = () => { const api = useBackendAPI(); @@ -21,6 +22,12 @@ export const useFlowRealtime = () => { const setIsGraphRunning = useGraphStore( useShallow((state) => state.setIsGraphRunning), ); + const updateEdgeBeads = useEdgeStore( + useShallow((state) => state.updateEdgeBeads), + ); + const resetEdgeBeads = useEdgeStore( + useShallow((state) => state.resetEdgeBeads), + ); const [{ flowExecutionID, flowID }] = useQueryStates({ flowExecutionID: parseAsString, @@ -34,12 +41,12 @@ export const useFlowRealtime = () => { if (data.graph_exec_id != flowExecutionID) { return; } - // TODO: Update the states of nodes updateNodeExecutionResult( data.node_id, data as unknown as NodeExecutionResult, ); updateStatus(data.node_id, data.status); + updateEdgeBeads(data.node_id, data as unknown as NodeExecutionResult); }, ); @@ -82,8 +89,9 @@ export const useFlowRealtime = () => { deregisterNodeExecutionEvent(); deregisterGraphExecutionSubscription(); deregisterGraphExecutionStatusEvent(); + resetEdgeBeads(); }; - }, [api, flowExecutionID]); + }, [api, flowExecutionID, resetEdgeBeads]); return {}; }; diff --git a/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/edges/CustomEdge.tsx b/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/edges/CustomEdge.tsx index 5887d8d012..b49fd11602 100644 --- a/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/edges/CustomEdge.tsx +++ b/autogpt_platform/frontend/src/app/(platform)/build/components/FlowEditor/edges/CustomEdge.tsx @@ -1,17 +1,30 @@ import { Button } from "@/components/atoms/Button/Button"; import { BaseEdge, + Edge as XYEdge, EdgeLabelRenderer, EdgeProps, getBezierPath, } from "@xyflow/react"; - import { useEdgeStore } from "@/app/(platform)/build/stores/edgeStore"; import { XIcon } from "@phosphor-icons/react"; +import { cn } from "@/lib/utils"; +import { NodeExecutionResult } from "@/lib/autogpt-server-api"; +import { JSBeads } from "./components/JSBeads"; + +export type CustomEdgeData = { + isStatic?: boolean; + beadUp?: number; + beadDown?: number; + beadData?: Map; +}; + +export type CustomEdge = XYEdge; import { memo } from "react"; const CustomEdge = ({ id, + data, sourceX, sourceY, targetX, @@ -20,8 +33,8 @@ const CustomEdge = ({ targetPosition, markerEnd, selected, -}: EdgeProps) => { - const removeConnection = useEdgeStore((state) => state.removeConnection); +}: EdgeProps) => { + const removeConnection = useEdgeStore((state) => state.removeEdge); const [edgePath, labelX, labelY] = getBezierPath({ sourceX, sourceY, @@ -31,14 +44,27 @@ const CustomEdge = ({ targetPosition, }); + const isStatic = data?.isStatic ?? false; + const beadUp = data?.beadUp ?? 0; + const beadDown = data?.beadDown ?? 0; + return ( <> +