From 379d38291435ac24ea4728dfe5d536eea9caa962 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Wed, 16 Apr 2025 15:50:20 +0800
Subject: [PATCH] feat: add current vars store

---
 .../workflow/current-vars-store/provider.tsx  | 30 ++++++
 .../workflow/current-vars-store/store.ts      | 91 +++++++++++++++++++
 web/app/components/workflow/index.tsx         | 13 ++-
 .../components/workflow/nodes/llm/panel.tsx   |  4 +-
 4 files changed, 132 insertions(+), 6 deletions(-)
 create mode 100644 web/app/components/workflow/current-vars-store/provider.tsx
 create mode 100644 web/app/components/workflow/current-vars-store/store.ts

diff --git a/web/app/components/workflow/current-vars-store/provider.tsx b/web/app/components/workflow/current-vars-store/provider.tsx
new file mode 100644
index 0000000000..154e7ff314
--- /dev/null
+++ b/web/app/components/workflow/current-vars-store/provider.tsx
@@ -0,0 +1,30 @@
+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
diff --git a/web/app/components/workflow/current-vars-store/store.ts b/web/app/components/workflow/current-vars-store/store.ts
new file mode 100644
index 0000000000..f8b40d9f07
--- /dev/null
+++ b/web/app/components/workflow/current-vars-store/store.ts
@@ -0,0 +1,91 @@
+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
+}
+
+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)
+    },
+  }))
+}
+
+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)
+}
diff --git a/web/app/components/workflow/index.tsx b/web/app/components/workflow/index.tsx
index 4c48afb56c..f748aceff5 100644
--- a/web/app/components/workflow/index.tsx
+++ b/web/app/components/workflow/index.tsx
@@ -102,6 +102,7 @@ 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 CurrentVarsProvider from './current-vars-store/provider'
 
 const nodeTypes = {
   [CUSTOM_NODE]: CustomNode,
@@ -453,11 +454,13 @@ const WorkflowWrap = memo(() => {
         edges={edgesData} >
         <FeaturesProvider features={initialFeatures}>
           <DatasetsDetailProvider nodes={nodesData}>
-            <Workflow
-              nodes={nodesData}
-              edges={edgesData}
-              viewport={data?.graph.viewport}
-            />
+            <CurrentVarsProvider>
+              <Workflow
+                nodes={nodesData}
+                edges={edgesData}
+                viewport={data?.graph.viewport}
+              />
+            </CurrentVarsProvider>
           </DatasetsDetailProvider>
         </FeaturesProvider>
       </WorkflowHistoryProvider>
diff --git a/web/app/components/workflow/nodes/llm/panel.tsx b/web/app/components/workflow/nodes/llm/panel.tsx
index 58b036317d..edffa65047 100644
--- a/web/app/components/workflow/nodes/llm/panel.tsx
+++ b/web/app/components/workflow/nodes/llm/panel.tsx
@@ -20,6 +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'
 
 const i18nPrefix = 'workflow.nodes.llm'
 
@@ -28,7 +29,8 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({
   data,
 }) => {
   const { t } = useTranslation()
-
+  const currentVars = useCurrentVarsStore(state => state.getVars())
+  console.log(currentVars)
   const {
     readOnly,
     inputs,