⚗️ add NodesProvider and context for managing nodes and edges

frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/ERDContent.tsx

@@ -6,14 +6,21 @@ import {
   Background,
   BackgroundVariant,
   type Edge,
+  type EdgeChange,
   type Node,
+  type NodeChange,
   type NodeMouseHandler,
   type OnNodeDrag,
   ReactFlow,
+  applyEdgeChanges,
+  applyNodeChanges,
   useEdgesState,
-  useNodesState,
 } from '@xyflow/react'
 import { type FC, useCallback } from 'react'
+import {
+  NodesProvider,
+  useNodesContext,
+} from '../../../providers/NodesProvider'
 import styles from './ERDContent.module.css'
 import { ERDContentProvider, useERDContentContext } from './ERDContentContext'
 import { NonRelatedTableGroupNode } from './NonRelatedTableGroupNode'
@@ -36,8 +43,6 @@ const edgeTypes = {
 }
 
 type Props = {
-  nodes: Node[]
-  edges: Edge[]
   enabledFeatures?:
     | {
         fitViewWhenActiveTableChange?: boolean | undefined
@@ -46,28 +51,21 @@ type Props = {
     | undefined
 }
 
-export const ERDContentInner: FC<Props> = ({
-  nodes: _nodes,
-  edges: _edges,
-  enabledFeatures,
-}) => {
-  const [nodes, setNodes, onNodesChange] = useNodesState<Node>(_nodes)
-  const [edges, setEdges, onEdgesChange] = useEdgesState<Edge>(_edges)
+export const ERDContentInner: FC<Props> = ({ enabledFeatures }) => {
+  const { nodes, setNodes, edges, setEdges } = useNodesContext()
   const {
     state: { loading },
   } = useERDContentContext()
   const { tableName: activeTableName } = useUserEditingActiveStore()
 
-  useInitialAutoLayout(
-    nodes,
-    enabledFeatures?.initialFitViewToActiveTable ?? true,
-  )
+  useInitialAutoLayout(enabledFeatures?.initialFitViewToActiveTable ?? true)
   useFitViewWhenActiveTableChange(
     enabledFeatures?.fitViewWhenActiveTableChange ?? true,
   )
   useSyncHighlightsActiveTableChange()
   useSyncHiddenNodesChange()
 
+  console.log(nodes)
   const { version } = useVersion()
   const handleNodeClick = useCallback(
     (tableId: string) => {
@@ -95,13 +93,20 @@ export const ERDContentInner: FC<Props> = ({
           hoverTableName: id,
         })
 
-      setEdges(updatedEdges)
-      setNodes(updatedNodes)
+      setEdges({
+        type: 'UPDATE_EDGES',
+        payload: updatedEdges,
+      })
+      setNodes({
+        type: 'UPDATE_DATA',
+        payload: updatedNodes,
+      })
     },
     [edges, nodes, setNodes, setEdges, activeTableName],
   )
 
   const handleMouseLeaveNode: NodeMouseHandler<Node> = useCallback(() => {
+    console.log('handleMouseLeaveNode')
     const { nodes: updatedNodes, edges: updatedEdges } = highlightNodesAndEdges(
       nodes,
       edges,
@@ -111,8 +116,14 @@ export const ERDContentInner: FC<Props> = ({
       },
     )
 
-    setEdges(updatedEdges)
-    setNodes(updatedNodes)
+    setEdges({
+      type: 'UPDATE_EDGES',
+      payload: updatedEdges,
+    })
+    setNodes({
+      type: 'UPDATE_DATA',
+      payload: updatedNodes,
+    })
   }, [edges, nodes, setNodes, setEdges, activeTableName])
 
   const handleDragStopNode: OnNodeDrag<Node> = useCallback(
@@ -132,6 +143,26 @@ export const ERDContentInner: FC<Props> = ({
     [version],
   )
 
+  const onNodesChange = useCallback(
+    (changes: NodeChange<Node>[]) => {
+      setNodes({
+        type: 'UPDATE_NODES',
+        payload: applyNodeChanges(changes, nodes),
+      })
+    },
+    [setNodes, nodes],
+  )
+
+  const onEdgesChange = useCallback(
+    (changes: EdgeChange<Edge>[]) => {
+      setEdges({
+        type: 'UPDATE_EDGES',
+        payload: applyEdgeChanges(changes, edges),
+      })
+    },
+    [setEdges, edges],
+  )
+
   const panOnDrag = [1, 2]
 
   return (

frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/ERDContentContext.tsx

@@ -6,14 +6,19 @@ import {
   useState,
 } from 'react'
 
+type LayoutStatus = 'not-started' | 'in-progress' | 'complete'
+
 type ERDContentContextState = {
   loading: boolean
   initializeComplete: boolean
+  //
+  // layoutStatus: LayoutStatus
 }
 
 type ERDContentContextActions = {
   setLoading: (loading: boolean) => void
   setInitializeComplete: (initializeComplete: boolean) => void
+  // setLayoutStatus: (layoutStatus: LayoutStatus) => void
 }
 
 type ERDContentConextValue = {
@@ -25,10 +30,13 @@ const ERDContentContext = createContext<ERDContentConextValue>({
   state: {
     loading: true,
     initializeComplete: false,
+    //
+    // layoutStatus: 'not-started',
   },
   actions: {
     setLoading: () => {},
     setInitializeComplete: () => {},
+    // setLayoutStatus: () => {},
   },
 })
 
@@ -37,12 +45,22 @@ export const useERDContentContext = () => useContext(ERDContentContext)
 export const ERDContentProvider: FC<PropsWithChildren> = ({ children }) => {
   const [loading, setLoading] = useState(true)
   const [initializeComplete, setInitializeComplete] = useState(false)
+  //
+  const [layoutStatus, setLayoutStatus] = useState<LayoutStatus>('not-started')
 
   return (
     <ERDContentContext.Provider
       value={{
-        state: { loading, initializeComplete },
-        actions: { setLoading, setInitializeComplete },
+        state: {
+          loading,
+          initializeComplete,
+          // layoutStatus,
+        },
+        actions: {
+          setLoading,
+          setInitializeComplete,
+          // setLayoutStatus,
+        },
       }}
     >
       {children}

frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/TableNode/TableDetail/RelatedTables/RelatedTables.tsx

@@ -1,14 +1,14 @@
 import { convertDBStructureToNodes } from '@/components/ERDRenderer/convertDBStructureToNodes'
 import { openRelatedTablesLogEvent } from '@/features/gtm/utils'
-import { useVersion } from '@/providers'
+import { NodesProvider, useNodesContext, useVersion } from '@/providers'
 import {
   replaceHiddenNodeIds,
   updateActiveTableName,
   useDBStructureStore,
 } from '@/stores'
 import type { Table } from '@liam-hq/db-structure'
 import { GotoIcon, IconButton } from '@liam-hq/ui'
-import { ReactFlowProvider, useReactFlow } from '@xyflow/react'
+import { ReactFlowProvider } from '@xyflow/react'
 import { type FC, useCallback } from 'react'
 import { ERDContent } from '../../../ERDContent'
 import styles from './RelatedTables.module.css'
@@ -25,11 +25,10 @@ export const RelatedTables: FC<Props> = ({ table }) => {
     dbStructure: extractedDBStructure,
     showMode: 'TABLE_NAME',
   })
-  const { getNodes } = useReactFlow()
+  const { nodes: mainPaneNodes } = useNodesContext()
   const { version } = useVersion()
   const handleClick = useCallback(() => {
     const visibleNodeIds: string[] = nodes.map((node) => node.id)
-    const mainPaneNodes = getNodes()
     const hiddenNodeIds = mainPaneNodes
       .filter((node) => !visibleNodeIds.includes(node.id))
       .map((node) => node.id)
@@ -42,7 +41,7 @@ export const RelatedTables: FC<Props> = ({ table }) => {
         cliVer: version.version,
         appEnv: version.envName,
       })
-  }, [nodes, getNodes, table.name, version])
+  }, [nodes, mainPaneNodes, table.name, version])
 
   return (
     <div className={styles.wrapper}>
@@ -55,16 +54,16 @@ export const RelatedTables: FC<Props> = ({ table }) => {
         />
       </div>
       <div className={styles.contentWrapper}>
-        <ReactFlowProvider>
-          <ERDContent
-            nodes={nodes}
-            edges={edges}
-            enabledFeatures={{
-              fitViewWhenActiveTableChange: false,
-              initialFitViewToActiveTable: false,
-            }}
-          />
-        </ReactFlowProvider>
+        <NodesProvider nodes={nodes} edges={edges}>
+          <ReactFlowProvider>
+            <ERDContent
+              enabledFeatures={{
+                fitViewWhenActiveTableChange: false,
+                initialFitViewToActiveTable: false,
+              }}
+            />
+          </ReactFlowProvider>
+        </NodesProvider>
       </div>
     </div>
   )

frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/Toolbar/TidyUpButton/TidyUpButton.tsx

@@ -3,15 +3,16 @@ import { useVersion } from '@/providers'
 import { useUserEditingStore } from '@/stores'
 import { IconButton, TidyUpIcon } from '@liam-hq/ui'
 import { ToolbarButton } from '@radix-ui/react-toolbar'
-import { useReactFlow } from '@xyflow/react'
 import { type FC, useCallback } from 'react'
+import { useNodesContext } from '../../../../../providers/NodesProvider'
 import { useAutoLayout } from '../../useAutoLayout'
 
 export const TidyUpButton: FC = () => {
-  const { getNodes, getEdges } = useReactFlow()
+  const { nodes, edges } = useNodesContext()
   const { handleLayout } = useAutoLayout()
   const { showMode } = useUserEditingStore()
   const { version } = useVersion()
+
   const handleClick = useCallback(() => {
     version.displayedOn === 'cli' &&
       toolbarActionLogEvent({
@@ -20,8 +21,8 @@ export const TidyUpButton: FC = () => {
         cliVer: version.version,
         appEnv: version.envName,
       })
-    handleLayout(getNodes(), getEdges())
-  }, [handleLayout, showMode, getNodes, getEdges, version])
+    handleLayout(nodes, edges)
+  }, [showMode, nodes, edges, version, handleLayout])
 
   return (
     <ToolbarButton asChild>

frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useAutoLayout/useAutoLayout.ts

@@ -1,11 +1,13 @@
+import { useNodesContext } from '@/providers'
 import { useReactFlow } from '@xyflow/react'
 import type { Edge, FitViewOptions, Node } from '@xyflow/react'
 import { useCallback } from 'react'
 import { useERDContentContext } from '../ERDContentContext'
 import { getElkLayout } from './getElkLayout'
 
 export const useAutoLayout = () => {
-  const { setNodes, fitView } = useReactFlow()
+  const { setNodes } = useNodesContext()
+  const { fitView } = useReactFlow()
   const {
     actions: { setLoading, setInitializeComplete },
   } = useERDContentContext()
@@ -38,14 +40,18 @@ export const useAutoLayout = () => {
         edges: visibleEdges,
       })
 
-      setNodes([...hiddenNodes, ...newNodes])
+      setNodes({
+        type: 'UPDATE_SIZE_AND_POSITION',
+        payload: [...hiddenNodes, ...newNodes],
+      })
+
       setTimeout(() => {
         fitView(fitViewOptions)
         setLoading(false)
         setInitializeComplete(true)
-      }, 0)
+      }, 1)
     },
-    [setNodes, fitView, setLoading, setInitializeComplete],
+    [setNodes, setLoading, setInitializeComplete, fitView],
   )
 
   return { handleLayout }

frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/useInitialAutoLayout.ts

@@ -1,7 +1,7 @@
+import { useNodesContext } from '@/providers'
 import type { QueryParam } from '@/schemas/queryParam'
 import { addHiddenNodeIds, updateActiveTableName } from '@/stores'
 import { decompressFromEncodedURIComponent } from '@/utils'
-import { type Node, useReactFlow } from '@xyflow/react'
 import { useEffect, useMemo } from 'react'
 import { useERDContentContext } from './ERDContentContext'
 import { highlightNodesAndEdges } from './highlightNodesAndEdges'
@@ -26,18 +26,16 @@ const getHiddenNodeIdsFromUrl = async (): Promise<string[]> => {
   return hiddenNodeIds ? hiddenNodeIds.split(',') : []
 }
 
-export const useInitialAutoLayout = (
-  nodes: Node[],
-  shouldFitViewToActiveTable: boolean,
-) => {
+export const useInitialAutoLayout = (shouldFitViewToActiveTable: boolean) => {
+  const { nodes, edges } = useNodesContext()
+
   const tableNodesInitialized = useMemo(
     () =>
       nodes
         .filter((node) => node.type === 'table')
         .some((node) => node.measured),
     [nodes],
   )
-  const { getEdges } = useReactFlow()
 
   const {
     state: { initializeComplete },
@@ -46,15 +44,14 @@ export const useInitialAutoLayout = (
 
   useEffect(() => {
     const initialize = async () => {
-      if (initializeComplete) {
+      if (initializeComplete || !tableNodesInitialized) {
         return
       }
 
       const activeTableName = getActiveTableNameFromUrl()
       updateActiveTableName(activeTableName)
       const hiddenNodeIds = await getHiddenNodeIdsFromUrl()
       addHiddenNodeIds(hiddenNodeIds)
-      const edges = getEdges()
       const hiddenNodes = nodes.map((node) => ({
         ...node,
         hidden: hiddenNodeIds.includes(node.id),
@@ -67,9 +64,7 @@ export const useInitialAutoLayout = (
           ? { maxZoom: 1, duration: 300, nodes: [{ id: activeTableName }] }
           : undefined
 
-      if (tableNodesInitialized) {
-        handleLayout(updatedNodes, updatedEdges, fitViewOptions)
-      }
+      handleLayout(updatedNodes, updatedEdges, fitViewOptions)
     }
 
     initialize()
@@ -78,7 +73,7 @@ export const useInitialAutoLayout = (
     initializeComplete,
     handleLayout,
     nodes,
-    getEdges,
+    edges,
     shouldFitViewToActiveTable,
   ])
 }