refactor: integrate currentNodeIdRef into data source store and update related components

web/app/components/datasets/documents/create-from-pipeline/data-source/online-documents/index.tsx

@@ -10,7 +10,7 @@ import { ssePost } from '@/service/base'
 import Toast from '@/app/components/base/toast'
 import type { DataSourceNodeCompletedResponse } from '@/types/pipeline'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
-import { useDataSourceStoreWithSelector } from '../store'
+import { useDataSourceStore, useDataSourceStoreWithSelector } from '../store'
 
 type OnlineDocumentsProps = {
   isInPipeline?: boolean
@@ -25,15 +25,11 @@ const OnlineDocuments = ({
 }: OnlineDocumentsProps) => {
   const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
   const documentsData = useDataSourceStoreWithSelector(state => state.documentsData)
-  const setDocumentsData = useDataSourceStoreWithSelector(state => state.setDocumentsData)
   const searchValue = useDataSourceStoreWithSelector(state => state.searchValue)
-  const setSearchValue = useDataSourceStoreWithSelector(state => state.setSearchValue)
-  const currentWorkspaceId = useDataSourceStoreWithSelector(state => state.currentWorkspaceId)
-  const setCurrentWorkspaceId = useDataSourceStoreWithSelector(state => state.setCurrentWorkspaceId)
-  const setOnlineDocuments = useDataSourceStoreWithSelector(state => state.setOnlineDocuments)
-  const setCurrentDocument = useDataSourceStoreWithSelector(state => state.setCurrentDocument)
   const selectedPagesId = useDataSourceStoreWithSelector(state => state.selectedPagesId)
-  const setSelectedPagesId = useDataSourceStoreWithSelector(state => state.setSelectedPagesId)
+  const currentWorkspaceId = useDataSourceStoreWithSelector(state => state.currentWorkspaceId)
+  const currentNodeIdRef = useDataSourceStoreWithSelector(state => state.currentNodeIdRef)
+  const dataSourceStore = useDataSourceStore()
 
   const PagesMapAndSelectedPagesId: DataSourceNotionPageMap = useMemo(() => {
     const pagesMap = (documentsData || []).reduce((prev: DataSourceNotionPageMap, next: DataSourceNotionWorkspace) => {
@@ -64,6 +60,7 @@ const OnlineDocuments = ({
       },
       {
         onDataSourceNodeCompleted: (documentsData: DataSourceNodeCompletedResponse) => {
+          const { setDocumentsData, setCurrentWorkspaceId } = dataSourceStore.getState()
           setDocumentsData(documentsData.data as DataSourceNotionWorkspace[])
           setCurrentWorkspaceId(documentsData.data[0].workspace_id)
         },
@@ -75,34 +72,58 @@ const OnlineDocuments = ({
         },
       },
     )
-  }, [datasourceNodeRunURL, setCurrentWorkspaceId, setDocumentsData])
+  }, [dataSourceStore, datasourceNodeRunURL])
 
   useEffect(() => {
-    if (!documentsData.length)
+    if (nodeId !== currentNodeIdRef.current) {
+      const {
+        setDocumentsData,
+        setCurrentWorkspaceId,
+        setSearchValue,
+        setSelectedPagesId,
+        setOnlineDocuments,
+        setCurrentDocument,
+      } = dataSourceStore.getState()
+      setDocumentsData([])
+      setCurrentWorkspaceId('')
+      setSearchValue('')
+      setSelectedPagesId(new Set())
+      setOnlineDocuments([])
+      setCurrentDocument(undefined)
+      currentNodeIdRef.current = nodeId
       getOnlineDocuments()
+    }
+    else {
+      // Avoid fetching documents when come back from next step
+      if (!documentsData.length)
+        getOnlineDocuments()
+    }
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [])
+  }, [nodeId])
 
   const currentWorkspace = documentsData.find(workspace => workspace.workspace_id === currentWorkspaceId)
 
   const handleSearchValueChange = useCallback((value: string) => {
+    const { setSearchValue } = dataSourceStore.getState()
     setSearchValue(value)
-  }, [setSearchValue])
+  }, [dataSourceStore])
 
   const handleSelectWorkspace = useCallback((workspaceId: string) => {
+    const { setCurrentWorkspaceId } = dataSourceStore.getState()
     setCurrentWorkspaceId(workspaceId)
-  }, [setCurrentWorkspaceId])
+  }, [dataSourceStore])
 
   const handleSelectPages = useCallback((newSelectedPagesId: Set<string>) => {
+    const { setSelectedPagesId, setOnlineDocuments } = dataSourceStore.getState()
     const selectedPages = Array.from(newSelectedPagesId).map(pageId => PagesMapAndSelectedPagesId[pageId])
-
     setSelectedPagesId(new Set(Array.from(newSelectedPagesId)))
     setOnlineDocuments(selectedPages)
-  }, [setSelectedPagesId, setOnlineDocuments, PagesMapAndSelectedPagesId])
+  }, [dataSourceStore, PagesMapAndSelectedPagesId])
 
   const handlePreviewPage = useCallback((previewPageId: string) => {
+    const { setCurrentDocument } = dataSourceStore.getState()
     setCurrentDocument(PagesMapAndSelectedPagesId[previewPageId])
-  }, [PagesMapAndSelectedPagesId, setCurrentDocument])
+  }, [PagesMapAndSelectedPagesId, dataSourceStore])
 
   const headerInfo = useMemo(() => {
     return {

web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/index.tsx

@@ -29,6 +29,7 @@ const OnlineDrive = ({
   const bucket = useDataSourceStoreWithSelector(state => state.bucket)
   const selectedFileList = useDataSourceStoreWithSelector(state => state.selectedFileList)
   const fileList = useDataSourceStoreWithSelector(state => state.fileList)
+  const currentNodeIdRef = useDataSourceStoreWithSelector(state => state.currentNodeIdRef)
   const dataSourceStore = useDataSourceStore()
   const [isLoading, setIsLoading] = useState(false)
 
@@ -82,10 +83,28 @@ const OnlineDrive = ({
   }, [prefix, bucket, datasourceNodeRunURL, dataSourceStore, fileList])
 
   useEffect(() => {
-    if (fileList.length > 0) return
-    getOnlineDriveFiles({})
+    if (nodeId !== currentNodeIdRef.current) {
+      const { setFileList, setBucket, setPrefix, setKeywords, setSelectedFileList } = dataSourceStore.getState()
+      setFileList([])
+      setBucket('')
+      setPrefix([])
+      setKeywords('')
+      setSelectedFileList([])
+      currentNodeIdRef.current = nodeId
+      getOnlineDriveFiles({
+        prefix: [],
+        bucket: '',
+        fileList: [],
+        startAfter: '',
+      })
+    }
+    else {
+      // Avoid fetching files when come back from next step
+      if (fileList.length > 0) return
+      getOnlineDriveFiles({})
+    }
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [])
+  }, [nodeId])
 
   const onlineDriveFileList = useMemo(() => {
     if (keywords)

web/app/components/datasets/documents/create-from-pipeline/data-source/store/index.ts

@@ -1,6 +1,8 @@
 import { useContext } from 'react'
 import { createStore, useStore } from 'zustand'
 import { DataSourceContext } from './provider'
+import type { CommonShape } from './slices/common'
+import { createCommonSlice } from './slices/common'
 import type { LocalFileSliceShape } from './slices/local-file'
 import { createLocalFileSlice } from './slices/local-file'
 import type { OnlineDocumentSliceShape } from './slices/online-document'
@@ -11,6 +13,7 @@ import type { OnlineDriveSliceShape } from './slices/online-drive'
 import { createOnlineDriveSlice } from './slices/online-drive'
 
 export type DataSourceShape =
+  CommonShape &
   LocalFileSliceShape &
   OnlineDocumentSliceShape &
   WebsiteCrawlSliceShape &
@@ -18,6 +21,7 @@ export type DataSourceShape =
 
 export const createDataSourceStore = () => {
   return createStore<DataSourceShape>((...args) => ({
+    ...createCommonSlice(...args),
     ...createLocalFileSlice(...args),
     ...createOnlineDocumentSlice(...args),
     ...createWebsiteCrawlSlice(...args),

web/app/components/datasets/documents/create-from-pipeline/data-source/store/slices/common.ts

@@ -0,0 +1,11 @@
+import type { StateCreator } from 'zustand'
+
+export type CommonShape = {
+  currentNodeIdRef: React.MutableRefObject<string | undefined>
+}
+
+export const createCommonSlice: StateCreator<CommonShape> = () => {
+  return ({
+    currentNodeIdRef: { current: undefined },
+  })
+}

web/app/components/datasets/documents/create-from-pipeline/data-source/website-crawl/index.tsx

@@ -20,7 +20,7 @@ import type {
   DataSourceNodeProcessingResponse,
 } from '@/types/pipeline'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
-import { useDataSourceStoreWithSelector } from '../store'
+import { useDataSourceStore, useDataSourceStoreWithSelector } from '../store'
 
 const I18N_PREFIX = 'datasetCreation.stepOne.website'
 
@@ -42,15 +42,10 @@ const WebsiteCrawl = ({
   const [crawlErrorMessage, setCrawlErrorMessage] = useState('')
   const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
   const crawlResult = useDataSourceStoreWithSelector(state => state.crawlResult)
-  const setCrawlResult = useDataSourceStoreWithSelector(state => state.setCrawlResult)
   const step = useDataSourceStoreWithSelector(state => state.step)
-  const setStep = useDataSourceStoreWithSelector(state => state.setStep)
   const checkedCrawlResult = useDataSourceStoreWithSelector(state => state.websitePages)
-  const setWebsitePages = useDataSourceStoreWithSelector(state => state.setWebsitePages)
-  const previewWebsitePageRef = useDataSourceStoreWithSelector(state => state.previewWebsitePageRef)
   const previewIndex = useDataSourceStoreWithSelector(state => state.previewIndex)
-  const setCurrentWebsite = useDataSourceStoreWithSelector(state => state.setCurrentWebsite)
-  const setPreviewIndex = useDataSourceStoreWithSelector(state => state.setPreviewIndex)
+  const dataSourceStore = useDataSourceStore()
 
   const usePreProcessingParams = useRef(!isInPipeline ? usePublishedPipelinePreProcessingParams : useDraftPipelinePreProcessingParams)
   const { data: paramsConfig, isFetching: isFetchingParams } = usePreProcessingParams.current({
@@ -63,6 +58,31 @@ const WebsiteCrawl = ({
       setControlFoldOptions(Date.now())
   }, [step])
 
+  useEffect(() => {
+    const {
+      setStep,
+      setCrawlResult,
+      setWebsitePages,
+      setPreviewIndex,
+      setCurrentWebsite,
+      currentNodeIdRef,
+      previewWebsitePageRef,
+    } = dataSourceStore.getState()
+    if (nodeId !== currentNodeIdRef.current) {
+      setStep(CrawlStep.init)
+      setCrawlResult(undefined)
+      setCurrentWebsite(undefined)
+      setWebsitePages([])
+      setPreviewIndex(0)
+      previewWebsitePageRef.current = undefined
+      setCrawledNum(0)
+      setTotalNum(0)
+      setCrawlErrorMessage('')
+      currentNodeIdRef.current = nodeId
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [nodeId])
+
   const isInit = step === CrawlStep.init
   const isCrawlFinished = step === CrawlStep.finished
   const isRunning = step === CrawlStep.running
@@ -72,16 +92,20 @@ const WebsiteCrawl = ({
     : `/rag/pipelines/${pipelineId}/workflows/draft/datasource/nodes/${nodeId}/run`
 
   const handleCheckedCrawlResultChange = useCallback((checkedCrawlResult: CrawlResultItem[]) => {
+    const { setWebsitePages, previewWebsitePageRef } = dataSourceStore.getState()
     setWebsitePages(checkedCrawlResult)
     previewWebsitePageRef.current = checkedCrawlResult[0]
-  }, [previewWebsitePageRef, setWebsitePages])
+  }, [dataSourceStore])
 
   const handlePreview = useCallback((website: CrawlResultItem, index: number) => {
+    const { setCurrentWebsite, setPreviewIndex } = dataSourceStore.getState()
     setCurrentWebsite(website)
     setPreviewIndex(index)
-  }, [setCurrentWebsite, setPreviewIndex])
+  }, [dataSourceStore])
 
   const handleRun = useCallback(async (value: Record<string, any>) => {
+    const { setStep, setCrawlResult } = dataSourceStore.getState()
+
     setStep(CrawlStep.running)
     ssePost(
       datasourceNodeRunURL,
@@ -120,7 +144,7 @@ const WebsiteCrawl = ({
         },
       },
     )
-  }, [datasourceNodeRunURL, handleCheckedCrawlResultChange, setCrawlResult, setStep, t])
+  }, [dataSourceStore, datasourceNodeRunURL, handleCheckedCrawlResultChange, t])
 
   const handleSubmit = useCallback((value: Record<string, any>) => {
     handleRun(value)