refactor(agentflow): Simplify ReactFlow and Context state sync (#5827)

j-sanaa · web-flow · commit 5cc657b900ac · 2026-02-23T16:55:09.000-08:00
* Fix biderectional sync

* addede utility function to sync states

* Fix comment

* Normalize nodes before updating states
diff --git a/packages/agentflow/src/Agentflow.tsx b/packages/agentflow/src/Agentflow.tsx
@@ -51,7 +51,7 @@ function AgentflowCanvas({
     renderHeader?: AgentflowProps['renderHeader']
     renderNodePalette?: AgentflowProps['renderNodePalette']
 }) {
-    const { state, setNodes, setEdges, setDirty, setReactFlowInstance, closeEditDialog } = useAgentflowContext()
+    const { state, setNodes, setEdges, setDirty, setReactFlowInstance, closeEditDialog, registerLocalStateSetters } = useAgentflowContext()
     const { isDarkMode } = useConfigContext()
     const agentflow = useAgentflow()
     const reactFlowWrapper = useRef<HTMLDivElement>(null)
@@ -75,71 +75,20 @@ function AgentflowCanvas({
     // Load available nodes
     const { availableNodes } = useFlowNodes()
 
-    // TODO: Refactor bidirectional state sync architecture
-    // Current implementation uses refs and setTimeout to prevent circular updates between
-    // ReactFlow's local state (useNodesState/useEdgesState) and AgentflowContext state.
-    // This approach is brittle and prone to race conditions.
-    //
-    // Problem: We have TWO sources of truth:
-    //   1. ReactFlow local state (nodes/edges) - updated by drag/drop, user interactions
-    //   2. AgentflowContext state - updated by edit dialog, delete, duplicate operations
-    //
-    // Current sync flow:
-    //   - Local ReactFlow → Context (for drag/drop) ✓
-    //   - Context → Local ReactFlow (for edit dialog updates) ⚠️ Brittle with refs/setTimeout
-    //
-    // Better solution (for future refactor):
-    //   Option A: Make context operations (updateNodeData, deleteNode, etc.) also call
-    //             setLocalNodes/setLocalEdges directly, then remove reverse sync entirely.
-    //   Option B: Use deep equality checks (e.g., lodash.isEqual) instead of refs/setTimeout.
-    //   Option C: Consolidate to single source of truth (context only, ReactFlow reads from it).
-    //
-    // See: https://github.com/FlowiseAI/Flowise/pull/XXX for discussion
-    //
-    // Use refs to prevent circular sync updates
-    const syncingToContext = useRef(false)
-    const syncingFromContext = useRef(false)
+    // Register local state setters with context on mount
+    useEffect(() => {
+        registerLocalStateSetters(setLocalNodes, setLocalEdges)
+    }, [registerLocalStateSetters, setLocalNodes, setLocalEdges])
 
     // Sync local ReactFlow state to context (when user interacts with canvas)
     useEffect(() => {
-        if (syncingFromContext.current) return
-        syncingToContext.current = true
         setNodes(nodes as FlowNode[])
-        // Use setTimeout to reset the flag after the sync completes
-        setTimeout(() => {
-            syncingToContext.current = false
-        }, 0)
     }, [nodes, setNodes])
 
     useEffect(() => {
-        if (syncingFromContext.current) return
-        syncingToContext.current = true
         setEdges(edges as FlowEdge[])
-        setTimeout(() => {
-            syncingToContext.current = false
-        }, 0)
     }, [edges, setEdges])
 
-    // Sync context state back to local ReactFlow (when updated via edit dialog, etc.)
-    useEffect(() => {
-        if (syncingToContext.current) return
-        syncingFromContext.current = true
-        setLocalNodes(state.nodes)
-        setTimeout(() => {
-            syncingFromContext.current = false
-        }, 0)
-    }, [state.nodes, setLocalNodes])
-
-    useEffect(() => {
-        if (syncingToContext.current) return
-        syncingFromContext.current = true
-        setLocalEdges(state.edges)
-        setTimeout(() => {
-            syncingFromContext.current = false
-        }, 0)
-    }, [state.edges, setLocalEdges])
-    // END TODO: Bidirectional state sync
-
     // Flow handlers
     const { handleConnect, handleNodesChange, handleEdgesChange, handleAddNode } = useFlowHandlers({
         nodes: nodes as FlowNode[],
diff --git a/packages/agentflow/src/infrastructure/store/AgentflowContext.tsx b/packages/agentflow/src/infrastructure/store/AgentflowContext.tsx
@@ -1,4 +1,4 @@
-import { createContext, Dispatch, ReactNode, useCallback, useContext, useReducer } from 'react'
+import { createContext, Dispatch, ReactNode, useCallback, useContext, useReducer, useRef } from 'react'
 import type { ReactFlowInstance } from 'reactflow'
 
 import type { AgentflowAction, AgentflowState, FlowConfig, FlowData, FlowEdge, FlowNode, InputParam, NodeData } from '@/core/types'
@@ -32,6 +32,9 @@ export interface AgentflowContextValue {
     //Dialog operations
     openEditDialog: (nodeId: string, data: NodeData, inputParams: InputParam[]) => void
     closeEditDialog: () => void
+
+    // Register ReactFlow local state setters
+    registerLocalStateSetters: (setLocalNodes: (nodes: FlowNode[]) => void, setLocalEdges: (edges: FlowEdge[]) => void) => void
 }
 
 const AgentflowContext = createContext<AgentflowContextValue | null>(null)
@@ -48,6 +51,34 @@ export function AgentflowStateProvider({ children, initialFlow }: AgentflowState
         edges: initialFlow?.edges || []
     })
 
+    // Store ReactFlow local state setters in refs which are populated by AgentflowCanvas
+    const localNodesSetterRef = useRef<((nodes: FlowNode[]) => void) | null>(null)
+    const localEdgesSetterRef = useRef<((edges: FlowEdge[]) => void) | null>(null)
+
+    const registerLocalStateSetters = useCallback(
+        (setLocalNodes: (nodes: FlowNode[]) => void, setLocalEdges: (edges: FlowEdge[]) => void) => {
+            localNodesSetterRef.current = setLocalNodes
+            localEdgesSetterRef.current = setLocalEdges
+        },
+        []
+    )
+
+    // Helper function to synchronize state updates between context and ReactFlow
+    const syncStateUpdate = useCallback(({ nodes, edges }: { nodes?: FlowNode[]; edges?: FlowEdge[] }) => {
+        if (nodes !== undefined) {
+            const normalizedNodes = normalizeNodes(nodes)
+            dispatch({ type: 'SET_NODES', payload: normalizedNodes })
+            localNodesSetterRef.current?.(normalizedNodes)
+        }
+        if (edges !== undefined) {
+            dispatch({ type: 'SET_EDGES', payload: edges })
+            localEdgesSetterRef.current?.(edges)
+        }
+        if (nodes !== undefined || edges !== undefined) {
+            dispatch({ type: 'SET_DIRTY', payload: true })
+        }
+    }, [])
+
     // Convenience setters
     const setNodes = useCallback((nodes: FlowNode[]) => {
         dispatch({ type: 'SET_NODES', payload: nodes })
@@ -74,11 +105,10 @@ export function AgentflowStateProvider({ children, initialFlow }: AgentflowState
         (nodeId: string) => {
             const newNodes = state.nodes.filter((node) => node.id !== nodeId)
             const newEdges = state.edges.filter((edge) => edge.source !== nodeId && edge.target !== nodeId)
-            dispatch({ type: 'SET_NODES', payload: newNodes })
-            dispatch({ type: 'SET_EDGES', payload: newEdges })
-            dispatch({ type: 'SET_DIRTY', payload: true })
+
+            syncStateUpdate({ nodes: newNodes, edges: newEdges })
         },
-        [state.nodes, state.edges]
+        [state.nodes, state.edges, syncStateUpdate]
     )
 
     const duplicateNode = useCallback(
@@ -96,10 +126,10 @@ export function AgentflowStateProvider({ children, initialFlow }: AgentflowState
                 data: { ...nodeToDuplicate.data }
             }
 
-            dispatch({ type: 'SET_NODES', payload: [...state.nodes, newNode] })
-            dispatch({ type: 'SET_DIRTY', payload: true })
+            const newNodes = [...state.nodes, newNode]
+            syncStateUpdate({ nodes: newNodes })
         },
-        [state.nodes]
+        [state.nodes, syncStateUpdate]
     )
 
     const updateNodeData = useCallback(
@@ -113,20 +143,19 @@ export function AgentflowStateProvider({ children, initialFlow }: AgentflowState
                 }
                 return node
             })
-            dispatch({ type: 'SET_NODES', payload: newNodes })
-            dispatch({ type: 'SET_DIRTY', payload: true })
+
+            syncStateUpdate({ nodes: newNodes })
         },
-        [state.nodes]
+        [state.nodes, syncStateUpdate]
     )
 
     // Edge operations
     const deleteEdge = useCallback(
         (edgeId: string) => {
             const newEdges = state.edges.filter((edge) => edge.id !== edgeId)
-            dispatch({ type: 'SET_EDGES', payload: newEdges })
-            dispatch({ type: 'SET_DIRTY', payload: true })
+            syncStateUpdate({ edges: newEdges })
         },
-        [state.edges]
+        [state.edges, syncStateUpdate]
     )
 
     // Dialog operations
@@ -179,7 +208,8 @@ export function AgentflowStateProvider({ children, initialFlow }: AgentflowState
         openEditDialog,
         closeEditDialog,
         getFlowData,
-        reset
+        reset,
+        registerLocalStateSetters
     }
 
     return <AgentflowContext.Provider value={value}>{children}</AgentflowContext.Provider>
