workflow store

web/app/components/workflow/context.tsx

-'use client'
-
-import { createContext, useContext } from 'use-context-selector'
-import type { Edge } from 'reactflow'
-import type {
-  BlockEnum,
-  Node,
-} from './types'
-
-export type WorkflowContextValue = {
-  mode: string
-  nodes: Node[]
-  edges: Edge[]
-  selectedNode?: Node
-  handleAddNextNode: (prevNode: Node, nextNodeType: BlockEnum) => void
-  handleUpdateNodeData: (nodeId: string, data: Node['data']) => void
-}
-
-export const WorkflowContext = createContext<WorkflowContextValue>({
-  mode: 'workflow',
-  nodes: [],
-  edges: [],
-  handleAddNextNode: () => {},
-  handleUpdateNodeData: () => {},
-})
-export const useWorkflowContext = () => useContext(WorkflowContext)

web/app/components/workflow/hooks.ts

-import type {
-  Dispatch,
-  SetStateAction,
-} from 'react'
 import { useCallback } from 'react'
 import produce from 'immer'
-import type {
-  Edge,
-  EdgeMouseHandler,
-} from 'reactflow'
-import type {
-  BlockEnum,
-  Node,
-} from './types'
-import { NodeInitialData } from './constants'
+import type { EdgeMouseHandler } from 'reactflow'
+import { useStoreApi } from 'reactflow'
+import type { SelectedNode } from './types'
+import { useStore } from './store'
 
-export const useWorkflow = (
-  nodes: Node[],
-  edges: Edge[],
-  setNodes: Dispatch<SetStateAction<Node[]>>,
-  setEdges: Dispatch<SetStateAction<Edge[]>>,
-) => {
-  const handleAddNextNode = useCallback((prevNode: Node, nextNodeType: BlockEnum) => {
-    const nextNode = {
-      id: `node-${Date.now()}`,
-      type: 'custom',
-      position: {
-        x: prevNode.position.x + 304,
-        y: prevNode.position.y,
-      },
-      data: NodeInitialData[nextNodeType],
-    }
-    const newEdge = {
-      id: `edge-${Date.now()}`,
-      source: prevNode.id,
-      target: nextNode.id,
-    }
-    setNodes((oldNodes) => {
-      return produce(oldNodes, (draft) => {
-        draft.push(nextNode)
-      })
-    })
-    setEdges((oldEdges) => {
-      return produce(oldEdges, (draft) => {
-        draft.push(newEdge)
-      })
-    })
-  }, [setNodes, setEdges])
+export const useWorkflow = () => {
+  const store = useStoreApi()
+  const setSelectedNode = useStore(state => state.setSelectedNode)
 
-  const handleUpdateNodeData = useCallback((nodeId: string, data: Node['data']) => {
-    setNodes((oldNodes) => {
-      return produce(oldNodes, (draft) => {
-        const node = draft.find(node => node.id === nodeId)
-        if (node)
-          node.data = data
-      })
-    })
-  }, [setNodes])
   const handleEnterEdge = useCallback<EdgeMouseHandler>((_, edge) => {
-    setEdges((oldEdges) => {
-      return produce(oldEdges, (draft) => {
-        const currentEdge = draft.find(e => e.id === edge.id)
-        if (currentEdge)
-          currentEdge.data = { ...currentEdge.data, hovering: true }
-      })
+    const {
+      edges,
+      setEdges,
+    } = store.getState()
+    const newEdges = produce(edges, (draft) => {
+      const currentEdge = draft.find(e => e.id === edge.id)
+      if (currentEdge)
+        currentEdge.data = { ...currentEdge.data, hovering: true }
     })
-  }, [setEdges])
+    setEdges(newEdges)
+  }, [store])
   const handleLeaveEdge = useCallback<EdgeMouseHandler>((_, edge) => {
-    setEdges((oldEdges) => {
-      return produce(oldEdges, (draft) => {
-        const currentEdge = draft.find(e => e.id === edge.id)
-        if (currentEdge)
-          currentEdge.data = { ...currentEdge.data, hovering: false }
-      })
+    const {
+      edges,
+      setEdges,
+    } = store.getState()
+    const newEdges = produce(edges, (draft) => {
+      const currentEdge = draft.find(e => e.id === edge.id)
+      if (currentEdge)
+        currentEdge.data = { ...currentEdge.data, hovering: false }
     })
-  }, [setEdges])
+    setEdges(newEdges)
+  }, [store])
+  const handleSelectNode = useCallback((selectNode: SelectedNode, cancelSelection?: boolean) => {
+    const {
+      getNodes,
+      setNodes,
+    } = store.getState()
+    if (cancelSelection) {
+      setSelectedNode(null)
+      const newNodes = produce(getNodes(), (draft) => {
+        const currentNode = draft.find(n => n.id === selectNode.id)
+
+        if (currentNode)
+          currentNode.data = { ...currentNode.data, selected: false }
+      })
+      setNodes(newNodes)
+    }
+    else {
+      setSelectedNode(selectNode)
+      const newNodes = produce(getNodes(), (draft) => {
+        const currentNode = draft.find(n => n.id === selectNode.id)
+
+        if (currentNode)
+          currentNode.data = { ...currentNode.data, selected: true }
+      })
+      setNodes(newNodes)
+    }
+  }, [setSelectedNode, store])
 
   return {
-    handleAddNextNode,
-    handleUpdateNodeData,
     handleEnterEdge,
     handleLeaveEdge,
+    handleSelectNode,
   }
 }

web/app/components/workflow/index.tsx

@@ -8,11 +8,6 @@ import ReactFlow, {
   useNodesState,
 } from 'reactflow'
 import 'reactflow/dist/style.css'
-import './style.css'
-import {
-  WorkflowContext,
-  useWorkflowContext,
-} from './context'
 import { useWorkflow } from './hooks'
 import Header from './header'
 import CustomNode from './nodes'
@@ -20,7 +15,6 @@ import ZoomInOut from './zoom-in-out'
 import CustomEdge from './custom-edge'
 import Panel from './panel'
 import type { Node } from './types'
-import { useStore } from './store'
 
 const nodeTypes = {
   custom: CustomNode,
@@ -29,13 +23,33 @@ const edgeTypes = {
   custom: CustomEdge,
 }
 
-const Workflow = memo(() => {
+type WorkflowProps = {
+  selectedNodeId?: string
+  nodes: Node[]
+  edges: Edge[]
+}
+const Workflow: FC<WorkflowProps> = memo(({
+  nodes: initialNodes,
+  edges: initialEdges,
+  selectedNodeId: initialSelectedNodeId,
+}) => {
+  const [nodes] = useNodesState(initialNodes)
+  const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges)
+
   const {
-    nodes,
-    edges,
-  } = useWorkflowContext()
-  const handleEnterEdge = useStore(state => state.handleEnterEdge)
-  const handleLeaveEdge = useStore(state => state.handleLeaveEdge)
+    handleEnterEdge,
+    handleLeaveEdge,
+    handleSelectNode,
+  } = useWorkflow()
+
+  useEffect(() => {
+    if (initialSelectedNodeId) {
+      const initialSelectedNode = nodes.find(n => n.id === initialSelectedNodeId)
+
+      if (initialSelectedNode)
+        handleSelectNode({ id: initialSelectedNodeId, data: initialSelectedNode.data })
+    }
+  }, [initialSelectedNodeId])
 
   return (
     <div className='relative w-full h-full'>
@@ -47,8 +61,10 @@ const Workflow = memo(() => {
         edgeTypes={edgeTypes}
         nodes={nodes}
         edges={edges}
+        onEdgesChange={onEdgesChange}
         onEdgeMouseEnter={handleEnterEdge}
         onEdgeMouseLeave={handleLeaveEdge}
+        multiSelectionKeyCode={null}
       >
         <Background
           gap={[14, 14]}
@@ -58,74 +74,10 @@ const Workflow = memo(() => {
     </div>
   )
 })
-Workflow.displayName = 'Workflow'
 
-type WorkflowWrapProps = {
-  selectedNodeId?: string
-  nodes: Node[]
-  edges: Edge[]
-}
-const WorkflowWrap: FC<WorkflowWrapProps> = memo(({
-  nodes: initialNodes,
-  edges: initialEdges,
-  selectedNodeId: initialSelectedNodeId,
-}) => {
-  const [nodes, setNodes] = useNodesState(initialNodes)
-  const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges)
-
-  const {
-    handleAddNextNode,
-    handleUpdateNodeData,
-    handleEnterEdge,
-    handleLeaveEdge,
-  } = useWorkflow(
-    nodes,
-    edges,
-    setNodes,
-    setEdges,
-  )
-  const handleSelectedNodeId = useStore(state => state.handleSelectedNodeId)
-  useEffect(() => {
-    if (initialSelectedNodeId)
-      handleSelectedNodeId(initialSelectedNodeId)
-  }, [initialSelectedNodeId, handleSelectedNodeId])
-  // const handleEnterEdge = useStore(state => state.handleEnterEdge)
-  // const handleLeaveEdge = useStore(state => state.handleLeaveEdge)
-
-  return (
-    <WorkflowContext.Provider value={{
-      mode: 'workflow',
-      nodes,
-      edges,
-      handleAddNextNode,
-      handleUpdateNodeData,
-    }}>
-      <div className='relative w-full h-full'>
-        <Header />
-        <Panel />
-        <ZoomInOut />
-        <ReactFlow
-          nodeTypes={nodeTypes}
-          edgeTypes={edgeTypes}
-          nodes={nodes}
-          edges={edges}
-          onEdgesChange={onEdgesChange}
-          onEdgeMouseEnter={handleEnterEdge}
-          onEdgeMouseLeave={handleLeaveEdge}
-        >
-          <Background
-            gap={[14, 14]}
-            size={1}
-          />
-        </ReactFlow>
-      </div>
-    </WorkflowContext.Provider>
-  )
-})
-
-WorkflowWrap.displayName = 'WorkflowWrap'
+Workflow.displayName = 'Workflow'
 
-const WorkflowWrapWithReactFlowProvider: FC<WorkflowWrapProps> = ({
+const WorkflowWrap: FC<WorkflowProps> = ({
   selectedNodeId,
   nodes,
   edges,
@@ -133,7 +85,7 @@ const WorkflowWrapWithReactFlowProvider: FC<WorkflowWrapProps> = ({
   return (
     <ReactFlowProvider>
       {selectedNodeId}
-      <WorkflowWrap
+      <Workflow
         selectedNodeId={selectedNodeId}
         nodes={nodes}
         edges={edges}
@@ -142,4 +94,4 @@ const WorkflowWrapWithReactFlowProvider: FC<WorkflowWrapProps> = ({
   )
 }
 
-export default memo(WorkflowWrapWithReactFlowProvider)
+export default memo(WorkflowWrap)

web/app/components/workflow/nodes/_base/components/next-step.tsx

-import type { FC } from 'react'
 import {
   memo,
   useCallback,
-  useMemo,
 } from 'react'
-import { getOutgoers } from 'reactflow'
 import BlockIcon from '../../../block-icon'
 import type { Node } from '../../../types'
 import { BlockEnum } from '../../../types'
-import { useWorkflowContext } from '../../../context'
+import { useStore } from '../../../store'
 import BlockSelector from '../../../block-selector'
 import { Plus } from '@/app/components/base/icons/src/vender/line/general'
 import Button from '@/app/components/base/button'
 
-type NextStepProps = {
-  selectedNode: Node
-}
-const NextStep: FC<NextStepProps> = ({
-  selectedNode,
-}) => {
-  const {
-    nodes,
-    edges,
-  } = useWorkflowContext()
-  const outgoers = useMemo(() => {
-    return getOutgoers(selectedNode, nodes, edges)
-  }, [selectedNode, nodes, edges])
+const NextStep = () => {
+  const selectedNode = useStore(state => state.selectedNode)
+  const outgoers: Node[] = []
 
   const renderAddNextNodeTrigger = useCallback((open: boolean) => {
     return (
@@ -60,7 +47,7 @@ const NextStep: FC<NextStepProps> = ({
   return (
     <div className='flex py-1'>
       <div className='shrink-0 relative flex items-center justify-center w-9 h-9 bg-white rounded-lg border-[0.5px] border-gray-200 shadow-xs'>
-        <BlockIcon type={selectedNode.data.type} />
+        <BlockIcon type={selectedNode!.data.type} />
       </div>
       <div className='shrink-0 w-6'></div>
       <div className='grow'>
@@ -74,7 +61,7 @@ const NextStep: FC<NextStepProps> = ({
                 type={outgoer.data.type}
                 className='shrink-0 mr-1.5'
               />
-              <div className='grow'>{outgoer.data.name}</div>
+              <div className='grow'>{outgoer.data.title}</div>
               <BlockSelector
                 onSelect={() => {}}
                 placement='top-end'
@@ -89,7 +76,7 @@ const NextStep: FC<NextStepProps> = ({
           ))
         }
         {
-          (!outgoers.length || selectedNode.data.type === BlockEnum.IfElse) && (
+          (!outgoers.length || selectedNode!.data.type === BlockEnum.IfElse) && (
             <BlockSelector
               onSelect={() => {}}
               placement='top'

web/app/components/workflow/nodes/_base/node.tsx

@@ -5,13 +5,10 @@ import type {
 import {
   cloneElement,
   memo,
-  useMemo,
 } from 'react'
 import type { NodeProps } from 'reactflow'
-import { useNodes } from 'reactflow'
-import { useStore } from '../../store'
-import type { NodeData } from '../../types'
 import BlockIcon from '../../block-icon'
+import { useWorkflow } from '../../hooks'
 import BlockSelector from '../../block-selector'
 import NodeControl from './components/node-control'
 
@@ -25,27 +22,22 @@ const BaseNode: FC<BaseNodeProps> = ({
   selected,
   children,
 }) => {
-  const nodes = useNodes<NodeData>()
-  const selectedNodeId = useStore(state => state.selectedNodeId)
-  const handleSelectedNodeId = useStore(state => state.handleSelectedNodeId)
-  const currentNode = useMemo(() => {
-    return nodes.find(node => node.id === nodeId)
-  }, [nodeId, nodes])
+  const { handleSelectNode } = useWorkflow()
 
   return (
     <div
       className={`
         group relative pb-2 w-[240px] bg-[#fcfdff] rounded-2xl shadow-xs
         hover:shadow-lg
-        ${selectedNodeId === nodeId ? 'border-[2px] border-primary-600' : 'border border-white'}
+        ${(data.selected && selected) ? 'border-[2px] border-primary-600' : 'border border-white'}
       `}
-      onClick={() => handleSelectedNodeId(nodeId || '')}
+      onClick={() => handleSelectNode({ id: nodeId, data })}
     >
       <NodeControl />
       <div className='flex items-center px-3 pt-3 pb-2'>
         <BlockIcon
           className='mr-2'
-          type={currentNode!.data.type}
+          type={data.type}
           size='md'
         />
         <div className='text-[13px] font-semibold text-gray-700'>

web/app/components/workflow/nodes/_base/panel.tsx

@@ -5,12 +5,10 @@ import type {
 import {
   cloneElement,
   memo,
-  useMemo,
 } from 'react'
-import type { NodeProps } from 'reactflow'
-import { useWorkflowContext } from '../../context'
-import { useStore } from '../../store'
+import type { SelectedNode } from '../../types'
 import BlockIcon from '../../block-icon'
+import { useWorkflow } from '../../hooks'
 import NextStep from './components/next-step'
 import {
   DotsHorizontal,
@@ -20,21 +18,14 @@ import { GitBranch01 } from '@/app/components/base/icons/src/vender/line/develop
 
 type BasePanelProps = {
   children: ReactElement
-} & Pick<NodeProps, 'id' | 'data'>
+} & SelectedNode
 
 const BasePanel: FC<BasePanelProps> = ({
   id,
   data,
   children,
 }) => {
-  const {
-    nodes,
-  } = useWorkflowContext()
-  const selectedNodeId = useStore(state => state.selectedNodeId)
-  const handleSelectedNodeId = useStore(state => state.handleSelectedNodeId)
-  const selectedNode = useMemo(() => {
-    return nodes.find(node => node.id === selectedNodeId)
-  }, [nodes, selectedNodeId])
+  const { handleSelectNode } = useWorkflow()
 
   return (
     <div className='mr-2 w-[420px] h-full bg-white shadow-lg border-[0.5px] border-gray-200 rounded-2xl z-10 overflow-y-auto'>
@@ -42,7 +33,7 @@ const BasePanel: FC<BasePanelProps> = ({
         <div className='flex items-center px-4 pt-3'>
           <BlockIcon
             className='shrink-0 mr-2'
-            type={selectedNode!.data.type}
+            type={data.type}
             size='md'
           />
           <div className='grow py-1 text-base text-gray-900 font-semibold '>{data.title}</div>
@@ -53,7 +44,7 @@ const BasePanel: FC<BasePanelProps> = ({
             <div className='mx-3 w-[1px] h-3.5 bg-gray-200' />
             <div
               className='flex items-center justify-center w-6 h-6 cursor-pointer'
-              onClick={() => handleSelectedNodeId('')}
+              onClick={() => handleSelectNode({ id, data }, true)}
             >
               <XClose className='w-4 h-4' />
             </div>
@@ -76,7 +67,7 @@ const BasePanel: FC<BasePanelProps> = ({
         <div className='mb-2 text-xs text-gray-400'>
           Add the next block in this workflow
         </div>
-        <NextStep selectedNode={selectedNode!} />
+        <NextStep />
       </div>
     </div>
   )

web/app/components/workflow/nodes/index.tsx

-import type { FC } from 'react'
 import { memo } from 'react'
 import type { NodeProps } from 'reactflow'
 import {
   Handle,
   Position,
 } from 'reactflow'
-import type { Node } from '../types'
+import type { SelectedNode } from '../types'
 import { BlockEnum } from '../types'
 import {
   NodeComponentMap,
@@ -14,7 +13,7 @@ import {
 import BaseNode from './_base/node'
 import BasePanel from './_base/panel'
 
-const CustomNode = (props: NodeProps) => {
+const CustomNode = memo((props: NodeProps) => {
   const nodeData = props.data
   const NodeComponent = NodeComponentMap[nodeData.type]
 
@@ -43,21 +42,15 @@ const CustomNode = (props: NodeProps) => {
       />
     </>
   )
-}
+})
+CustomNode.displayName = 'CustomNode'
 
-type PanelProps = {
-  node: Node
-}
-export const Panel: FC<PanelProps> = memo(({
-  node,
-}) => {
-  const PanelComponent = PanelComponentMap[node.data.type]
+export const Panel = memo((props: SelectedNode) => {
+  const nodeData = props.data
+  const PanelComponent = PanelComponentMap[nodeData.type]
 
   return (
-    <BasePanel
-      id={node.id}
-      data={node.data}
-    >
+    <BasePanel {...props}>
       <PanelComponent />
     </BasePanel>
   )

web/app/components/workflow/panel/index.tsx

@@ -3,22 +3,14 @@ import {
   memo,
   useMemo,
 } from 'react'
-import { useNodes } from 'reactflow'
-import { useWorkflowContext } from '../context'
 import { Panel as NodePanel } from '../nodes'
 import { useStore } from '../store'
 import WorkflowInfo from './workflow-info'
 import DebugAndPreview from './debug-and-preview'
 
 const Panel: FC = () => {
-  const {
-    mode,
-  } = useWorkflowContext()
-  const nodes = useNodes()
-  const selectedNodeId = useStore(state => state.selectedNodeId)
-  const selectedNode = useMemo(() => {
-    return nodes.find(node => node.id === selectedNodeId)
-  }, [nodes, selectedNodeId])
+  const mode = useStore(state => state.mode)
+  const selectedNode = useStore(state => state.selectedNode)
   const {
     showWorkflowInfoPanel,
     showNodePanel,
@@ -26,7 +18,7 @@ const Panel: FC = () => {
   } = useMemo(() => {
     return {
       showWorkflowInfoPanel: mode === 'workflow' && !selectedNode,
-      showNodePanel: selectedNode,
+      showNodePanel: !!selectedNode,
       showDebugAndPreviewPanel: mode === 'chatbot' && !selectedNode,
     }
   }, [mode, selectedNode])
@@ -35,7 +27,7 @@ const Panel: FC = () => {
     <div className='absolute top-14 right-0 bottom-2 flex'>
       {
         showNodePanel && (
-          <NodePanel node={selectedNode as any} />
+          <NodePanel {...selectedNode!} />
         )
       }
       {

web/app/components/workflow/store.ts

 import { create } from 'zustand'
-import type { EdgeMouseHandler } from 'reactflow'
+import type { SelectedNode } from './types'
 
 type State = {
   mode: string
-  selectedNodeId: string
-  hoveringEdgeId: string
+  selectedNode: SelectedNode | null
 }
 
 type Action = {
-  handleSelectedNodeId: (selectedNodeId: State['selectedNodeId']) => void
-  handleEnterEdge: EdgeMouseHandler
-  handleLeaveEdge: EdgeMouseHandler
+  setSelectedNode: (node: SelectedNode | null) => void
 }
 
 export const useStore = create<State & Action>(set => ({
   mode: 'workflow',
-  selectedNodeId: '',
-  handleSelectedNodeId: selectedNodeId => set(() => ({ selectedNodeId })),
-  hoveringEdgeId: '',
-  handleEnterEdge: (_, edge) => set(() => ({ hoveringEdgeId: edge.id })),
-  handleLeaveEdge: () => set(() => ({ hoveringEdgeId: '' })),
+  selectedNode: null,
+  setSelectedNode: node => set(() => ({ selectedNode: node })),
 }))

web/app/components/workflow/style.css

-.react-flow__edge-path {
-  stroke-width: 2px !important;
-}
\ No newline at end of file

web/app/components/workflow/types.ts

@@ -15,20 +15,16 @@ export enum BlockEnum {
   Tool = 'tool',
 }
 
-export type NodeData = {
-  type: BlockEnum
-  name?: string
-  icon?: any
-  description?: string
-}
-export type Node = ReactFlowNode<NodeData>
-
 export type CommonNodeType = {
   title: string
   desc: string
-  type: string
+  type: BlockEnum
+  selected?: boolean
 }
 
+export type Node = ReactFlowNode<CommonNodeType>
+export type SelectedNode = Pick<Node, 'id' | 'data'>
+
 export type ValueSelector = string[] // [nodeId, key | obj key path]
 
 export type Variable = {