chore: refactor to slice

web/app/components/workflow/current-vars-store/provider.tsx

@@ -1,30 +0,0 @@
-import type { FC } from 'react'
-import { createContext, useRef } from 'react'
-import { createCurrentVarsStore } from './store'
-
-type CurrentVarsStoreApi = ReturnType<typeof createCurrentVarsStore>
-
-type CurrentVarsContextType = CurrentVarsStoreApi | undefined
-
-export const CurrentVarsContext = createContext<CurrentVarsContextType>(undefined)
-
-type CurrentVarsProviderProps = {
-  children: React.ReactNode
-}
-
-const CurrentVarsProvider: FC<CurrentVarsProviderProps> = ({
-  children,
-}) => {
-  const storeRef = useRef<CurrentVarsStoreApi>()
-
-  if (!storeRef.current)
-    storeRef.current = createCurrentVarsStore()
-
-  return (
-    <CurrentVarsContext.Provider value={storeRef.current!}>
-      {children}
-    </CurrentVarsContext.Provider>
-  )
-}
-
-export default CurrentVarsProvider

web/app/components/workflow/current-vars-store/store.ts

@@ -1,124 +0,0 @@
-import { useContext } from 'react'
-import { createStore, useStore } from 'zustand'
-import { CurrentVarsContext } from './provider'
-import produce from 'immer'
-
-type NodeVars = {
-  id: string
-  name: string
-  type: string
-  vars: {
-    key: string
-    type: string
-    value: any
-  }[]
-}
-
-type CurrentVarsState = {
-  nodes: NodeVars[]
-}
-
-type CurrentVarsActions = {
-  setVars: (vars: NodeVars[]) => void
-  getVars: () => NodeVars[]
-  clearVars: () => void
-  setNodeVars: (nodeId: string, payload: NodeVars) => void
-  clearNodeVars: (nodeId: string) => void
-  getNodeVars: (nodeId: string) => NodeVars | undefined
-  hasNodeVars: (nodeId: string) => boolean
-  setVar: (nodeId: string, key: string, value: any) => void
-  getVar: (nodeId: string, key: string) => any
-}
-
-type CurrentVarsStore = CurrentVarsState & CurrentVarsActions
-
-export const createCurrentVarsStore = () => {
-  return createStore<CurrentVarsStore>((set, get) => ({
-    nodes: [{
-      id: '',
-      name: '',
-      type: '',
-      vars: [],
-    }],
-    setVars: (vars) => {
-      set(() => ({
-        nodes: vars,
-      }))
-    },
-    getVars: () => {
-      return get().nodes
-    },
-    clearVars: () => {
-      set(() => ({
-        nodes: [],
-      }))
-    },
-    setNodeVars: (nodeId, vars) => {
-      set((state) => {
-        // eslint-disable-next-line sonarjs/no-nested-functions
-        const nodes = state.nodes.map((node) => {
-          if (node.id === nodeId) {
-            return produce(node, (draft) => {
-              draft.vars = vars.vars
-            })
-          }
-
-          return node
-        })
-        return {
-          nodes,
-        }
-      })
-    },
-    clearNodeVars: (nodeId) => {
-      set(produce((state: CurrentVarsStore) => {
-        // eslint-disable-next-line sonarjs/no-nested-functions
-        const nodes = state.nodes.filter(node => node.id !== nodeId)
-        state.nodes = nodes
-      },
-      ))
-    },
-    getNodeVars: (nodeId) => {
-      const nodes = get().nodes
-      return nodes.find(node => node.id === nodeId)
-    },
-    hasNodeVars: (nodeId) => {
-      return !!get().getNodeVars(nodeId)
-    },
-    setVar: (nodeId, key, value) => {
-      set(produce((state: CurrentVarsStore) => {
-        // eslint-disable-next-line sonarjs/no-nested-functions
-        const nodes = state.nodes.map((node) => {
-          if (node.id === nodeId) {
-            return produce(node, (draft) => {
-              const index = draft.vars.findIndex(v => v.key === key)
-              if (index !== -1)
-                draft.vars[index].value = value
-            })
-          }
-          return node
-        })
-        state.nodes = nodes
-      }))
-    },
-    getVar(nodeId, key) {
-      const node = get().getNodeVars(nodeId)
-      if (!node)
-        return undefined
-
-      const variable = node.vars.find(v => v.key === key)
-      if (!variable)
-        return undefined
-
-      return variable.value
-    },
-  }))
-}
-
-export const useCurrentVarsStore = <T>(selector: (state: CurrentVarsStore) => T): T => {
-  const store = useContext(CurrentVarsContext)
-  if (!store)
-    throw new Error('Missing CurrentVarsContext.Provider in the tree')
-
-  return useStore(store, selector)
-}

web/app/components/workflow/hooks/use-current-vars.ts

@@ -1,12 +1,15 @@
-import { useCurrentVarsStore } from '../current-vars-store/store'
-import { useLastRunStore } from '../last-run-store/store'
+import { useWorkflowStore } from '../store'
 const useCurrentVars = () => {
-  const currentVars = useCurrentVarsStore(state => state.nodes)
-  const getCurrentVar = useCurrentVarsStore(state => state.getVar)
-  const getLastRunVar = useLastRunStore(state => state.getVar)
-  const setCurrentVar = useCurrentVarsStore(state => state.setVar)
-  const clearCurrentVars = useCurrentVarsStore(state => state.clearVars)
-  const clearNodeVars = useCurrentVarsStore(state => state.clearNodeVars)
+  const workflowStore = useWorkflowStore()
+  const {
+    currentNodes,
+    getCurrentVar,
+    setCurrentVar,
+    clearCurrentVars,
+    clearCurrentNodeVars,
+    getLastRunVar,
+    getLastRunInfos,
+  } = workflowStore.getState()
 
   const isVarChanged = (nodeId: string, key: string) => {
     return getCurrentVar(nodeId, key) !== getLastRunVar(nodeId, key)
@@ -19,10 +22,11 @@ const useCurrentVars = () => {
   }
 
   return {
-    currentVars,
+    currentVars: currentNodes,
+    getLastRunInfos,
     isVarChanged,
     clearCurrentVars,
-    clearNodeVars,
+    clearCurrentNodeVars,
     setCurrentVar,
     resetToLastRunVar,
   }

web/app/components/workflow/index.tsx

@@ -102,8 +102,6 @@ import Confirm from '@/app/components/base/confirm'
 import { FILE_EXTS } from '@/app/components/base/prompt-editor/constants'
 import { fetchFileUploadConfig } from '@/service/common'
 import DatasetsDetailProvider from './datasets-detail-store/provider'
-import LastRunProvider from './last-run-store/provider'
-import CurrentVarsProvider from './current-vars-store/provider'
 
 const nodeTypes = {
   [CUSTOM_NODE]: CustomNode,
@@ -455,15 +453,11 @@ const WorkflowWrap = memo(() => {
         edges={edgesData} >
         <FeaturesProvider features={initialFeatures}>
           <DatasetsDetailProvider nodes={nodesData}>
-            <LastRunProvider>
-              <CurrentVarsProvider>
-                <Workflow
-                  nodes={nodesData}
-                  edges={edgesData}
-                  viewport={data?.graph.viewport}
-                />
-              </CurrentVarsProvider>
-            </LastRunProvider>
+            <Workflow
+              nodes={nodesData}
+              edges={edgesData}
+              viewport={data?.graph.viewport}
+            />
           </DatasetsDetailProvider>
         </FeaturesProvider>
       </WorkflowHistoryProvider>

web/app/components/workflow/last-run-store/provider.tsx

@@ -1,30 +0,0 @@
-import type { FC } from 'react'
-import { createContext, useRef } from 'react'
-import { createLastRunStore } from './store'
-
-type LastRunStoreApi = ReturnType<typeof createLastRunStore>
-
-type LastRunContextType = LastRunStoreApi | undefined
-
-export const LastRunContext = createContext<LastRunContextType>(undefined)
-
-type LastRunProviderProps = {
-  children: React.ReactNode
-}
-
-const LastRunProvider: FC<LastRunProviderProps> = ({
-  children,
-}) => {
-  const storeRef = useRef<LastRunStoreApi>()
-
-  if (!storeRef.current)
-    storeRef.current = createLastRunStore()
-
-  return (
-    <LastRunContext.Provider value={storeRef.current!}>
-      {children}
-    </LastRunContext.Provider>
-  )
-}
-
-export default LastRunProvider

web/app/components/workflow/last-run-store/store.ts

@@ -1,79 +0,0 @@
-import { useContext } from 'react'
-import { createStore, useStore } from 'zustand'
-import { LastRunContext } from './provider'
-
-type NodeInfo = {
-  id: string
-  name: string
-  type: string
-  vars: {
-    key: string
-    type: string
-    value: any
-  }[]
-} & {
-  input: Record<string, any>
-  output: Record<string, any>
-}
-
-type LastRunState = {
-  nodes: NodeInfo[]
-}
-
-type LastRunActions = {
-  setInfos: (vars: NodeInfo[]) => void
-  getInfos: () => NodeInfo[]
-  getNodeInfo: (nodeId: string) => NodeInfo | undefined
-  getVar: (nodeId: string, key: string) => any
-}
-
-type LastRunStore = LastRunState & LastRunActions
-
-export const createLastRunStore = () => {
-  return createStore<LastRunStore>((set, get) => ({
-    nodes: [{
-      id: '',
-      name: '',
-      type: '',
-      vars: [],
-      input: {},
-      output: {},
-    }],
-    setInfos: (vars) => {
-      set(() => ({
-        nodes: vars,
-      }))
-    },
-    getInfos: () => {
-      return get().nodes
-    },
-    clearVars: () => {
-      set(() => ({
-        nodes: [],
-      }))
-    },
-    getNodeInfo: (nodeId) => {
-      const nodes = get().nodes
-      return nodes.find(node => node.id === nodeId)
-    },
-    getVar: (nodeId, key) => {
-      const node = get().getNodeInfo(nodeId)
-      if (!node)
-        return undefined
-
-      const varItem = node.vars.find(v => v.key === key)
-      if (!varItem)
-        return undefined
-
-      return varItem.value
-    },
-  }))
-}
-
-export const useLastRunStore = <T>(selector: (state: LastRunStore) => T): T => {
-  const store = useContext(LastRunContext)
-  if (!store)
-    throw new Error('Missing LastRunContext.Provider in the tree')
-
-  return useStore(store, selector)
-}

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

@@ -20,8 +20,7 @@ import type { Props as FormProps } from '@/app/components/workflow/nodes/_base/c
 import ResultPanel from '@/app/components/workflow/run/result-panel'
 import Tooltip from '@/app/components/base/tooltip'
 import Editor from '@/app/components/workflow/nodes/_base/components/prompt/editor'
-import { useCurrentVarsStore } from '../../current-vars-store/store'
-import { useLastRunStore } from '../../last-run-store/store'
+import useCurrentVars from '../../hooks/use-current-vars'
 
 const i18nPrefix = 'workflow.nodes.llm'
 
@@ -30,9 +29,11 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({
   data,
 }) => {
   const { t } = useTranslation()
-  const currentVars = useCurrentVarsStore(state => state.getVars())
-  const lastRunInfo = useLastRunStore(state => state.getInfos())
-  console.log(currentVars, lastRunInfo)
+  const {
+    currentVars,
+    getLastRunInfos,
+  } = useCurrentVars()
+  console.log(currentVars, getLastRunInfos())
   const {
     readOnly,
     inputs,

web/app/components/workflow/store/workflow/current-vars-slice.ts

@@ -0,0 +1,112 @@
+import type { StateCreator } from 'zustand'
+
+import produce from 'immer'
+
+type NodeVars = {
+  id: string
+  name: string
+  type: string
+  vars: {
+    key: string
+    type: string
+    value: any
+  }[]
+}
+
+type CurrentVarsState = {
+  currentNodes: NodeVars[]
+}
+
+type CurrentVarsActions = {
+  setCurrentVars: (vars: NodeVars[]) => void
+  getCurrentVars: () => NodeVars[]
+  clearCurrentVars: () => void
+  setCurrentNodeVars: (nodeId: string, payload: NodeVars) => void
+  clearCurrentNodeVars: (nodeId: string) => void
+  getCurrentNodeVars: (nodeId: string) => NodeVars | undefined
+  hasCurrentNodeVars: (nodeId: string) => boolean
+  setCurrentVar: (nodeId: string, key: string, value: any) => void
+  getCurrentVar: (nodeId: string, key: string) => any
+}
+
+export type CurrentVarsSliceShape = CurrentVarsState & CurrentVarsActions
+
+export const createCurrentVarsSlice: StateCreator<CurrentVarsSliceShape> = (set, get) => {
+  return ({
+    currentNodes: [{
+      id: 'abc',
+      name: '',
+      type: '',
+      vars: [],
+    }],
+    setCurrentVars: (vars) => {
+      set(() => ({
+        currentNodes: vars,
+      }))
+    },
+    getCurrentVars: () => {
+      return get().currentNodes
+    },
+    clearCurrentVars: () => {
+      set(() => ({
+        currentNodes: [],
+      }))
+    },
+    setCurrentNodeVars: (nodeId, vars) => {
+      set((state) => {
+        const nodes = state.currentNodes.map((node) => {
+          if (node.id === nodeId) {
+            return produce(node, (draft) => {
+              draft.vars = vars.vars
+            })
+          }
+
+          return node
+        })
+        return {
+          currentNodes: nodes,
+        }
+      })
+    },
+    clearCurrentNodeVars: (nodeId) => {
+      set(produce((state: CurrentVarsSliceShape) => {
+        const nodes = state.currentNodes.filter(node => node.id !== nodeId)
+        state.currentNodes = nodes
+      },
+      ))
+    },
+    getCurrentNodeVars: (nodeId) => {
+      const nodes = get().currentNodes
+      return nodes.find(node => node.id === nodeId)
+    },
+    hasCurrentNodeVars: (nodeId) => {
+      return !!get().getCurrentNodeVars(nodeId)
+    },
+    setCurrentVar: (nodeId, key, value) => {
+      set(produce((state: CurrentVarsSliceShape) => {
+        const nodes = state.currentNodes.map((node) => {
+          if (node.id === nodeId) {
+            return produce(node, (draft) => {
+              const index = draft.vars.findIndex(v => v.key === key)
+              if (index !== -1)
+                draft.vars[index].value = value
+            })
+          }
+          return node
+        })
+        state.currentNodes = nodes
+      }))
+    },
+    getCurrentVar(nodeId, key) {
+      const node = get().getCurrentNodeVars(nodeId)
+      if (!node)
+        return undefined
+
+      const variable = node.vars.find(v => v.key === key)
+      if (!variable)
+        return undefined
+
+      return variable.value
+    },
+  })
+}

web/app/components/workflow/store/workflow/index.ts

@@ -25,6 +25,11 @@ import type { WorkflowDraftSliceShape } from './workflow-draft-slice'
 import { createWorkflowDraftSlice } from './workflow-draft-slice'
 import type { WorkflowSliceShape } from './workflow-slice'
 import { createWorkflowSlice } from './workflow-slice'
+import type { LastRunSliceShape } from './last-run-slice'
+import { createLastRunSlice } from './last-run-slice'
+import type { CurrentVarsSliceShape } from './current-vars-slice'
+import { createCurrentVarsSlice } from './current-vars-slice'
+
 import { WorkflowContext } from '@/app/components/workflow/context'
 
 export type Shape =
@@ -38,7 +43,9 @@ export type Shape =
   ToolSliceShape &
   VersionSliceShape &
   WorkflowDraftSliceShape &
-  WorkflowSliceShape
+  WorkflowSliceShape &
+  LastRunSliceShape &
+  CurrentVarsSliceShape
 
 export const createWorkflowStore = () => {
   return createStore<Shape>((...args) => ({
@@ -53,6 +60,8 @@ export const createWorkflowStore = () => {
     ...createVersionSlice(...args),
     ...createWorkflowDraftSlice(...args),
     ...createWorkflowSlice(...args),
+    ...createLastRunSlice(...args),
+    ...createCurrentVarsSlice(...args),
   }))
 }
 

web/app/components/workflow/store/workflow/last-run-slice.ts

@@ -0,0 +1,69 @@
+import type { StateCreator } from 'zustand'
+
+type NodeInfo = {
+  id: string
+  name: string
+  type: string
+  vars: {
+    key: string
+    type: string
+    value: any
+  }[]
+} & {
+  input: Record<string, any>
+  output: Record<string, any>
+}
+
+type LastRunState = {
+  nodes: NodeInfo[]
+}
+
+type LastRunActions = {
+  setLastRunInfos: (vars: NodeInfo[]) => void
+  getLastRunInfos: () => NodeInfo[]
+  getLastRunNodeInfo: (nodeId: string) => NodeInfo | undefined
+  getLastRunVar: (nodeId: string, key: string) => any
+}
+
+export type LastRunSliceShape = LastRunState & LastRunActions
+
+export const createLastRunSlice: StateCreator<LastRunSliceShape> = (set, get) => {
+  return ({
+    nodes: [{
+      id: 'test',
+      name: '',
+      type: '',
+      vars: [],
+      input: {},
+      output: {},
+    }],
+    setLastRunInfos: (vars) => {
+      set(() => ({
+        nodes: vars,
+      }))
+    },
+    getLastRunInfos: () => {
+      return get().nodes
+    },
+    clearVars: () => {
+      set(() => ({
+        nodes: [],
+      }))
+    },
+    getLastRunNodeInfo: (nodeId) => {
+      const nodes = get().nodes
+      return nodes.find(node => node.id === nodeId)
+    },
+    getLastRunVar: (nodeId, key) => {
+      const node = get().getLastRunNodeInfo(nodeId)
+      if (!node)
+        return undefined
+
+      const varItem = node.vars.find(v => v.key === key)
+      if (!varItem)
+        return undefined
+
+      return varItem.value
+    },
+  })
+}