workflow store

2024-02-23 14:15:43 +08:00 · 2024-02-23 14:15:43 +08:00 · f6c07c996b
parent 7ba0bfffa2
commit f6c07c996b
11 changed files with 127 additions and 275 deletions
--- a/web/app/components/workflow/context.tsx
+++ b/web/app/components/workflow/context.tsx
@ -1,26 +0,0 @@
-'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)
--- a/web/app/components/workflow/hooks.ts
+++ b/web/app/components/workflow/hooks.ts
@ -1,84 +1,68 @@
-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,
  }
 }
--- a/web/app/components/workflow/index.tsx
+++ b/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'
-
-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)
--- a/web/app/components/workflow/nodes/_base/components/next-step.tsx
+++ b/web/app/components/workflow/nodes/_base/components/next-step.tsx
@ -1,31 +1,18 @@
-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'
--- a/web/app/components/workflow/nodes/_base/node.tsx
+++ b/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'>
--- a/web/app/components/workflow/nodes/_base/panel.tsx
+++ b/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>
  )
--- a/web/app/components/workflow/nodes/index.tsx
+++ b/web/app/components/workflow/nodes/index.tsx
@ -1,11 +1,10 @@
-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>
  )
--- a/web/app/components/workflow/panel/index.tsx
+++ b/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!} />
        )
      }
      {
--- a/web/app/components/workflow/store.ts
+++ b/web/app/components/workflow/store.ts
@ -1,23 +1,17 @@
 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 })),
 }))
--- a/web/app/components/workflow/style.css
+++ b/web/app/components/workflow/style.css
@ -1,3 +0,0 @@
-.react-flow__edge-path {
-  stroke-width: 2px !important;
-}
--- a/web/app/components/workflow/types.ts
+++ b/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 = {