From b3431ab0c40f7619fabae753a91b93e5f6c5e513 Mon Sep 17 00:00:00 2001
From: twwu <twwu@dify.ai>
Date: Tue, 8 Jul 2025 16:46:59 +0800
Subject: [PATCH] feat: Refactor online drive components to improve file
 retrieval and selection logic

---
 .../header/breadcrumbs/dropdown/index.tsx     |  2 +-
 .../file-list/header/breadcrumbs/index.tsx    | 33 +++++++++++++++----
 .../online-drive/file-list/header/index.tsx   |  9 +++++
 .../online-drive/file-list/index.tsx          |  9 +++++
 .../online-drive/file-list/list/index.tsx     | 14 ++++++--
 .../data-source/online-drive/index.tsx        | 33 +++++++++++++------
 .../documents/create-from-pipeline/index.tsx  | 28 +++++++++-------
 web/app/components/datasets/list/datasets.tsx |  3 +-
 .../components/panel/test-run/hooks.ts        | 25 --------------
 9 files changed, 100 insertions(+), 56 deletions(-)

diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/breadcrumbs/dropdown/index.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/breadcrumbs/dropdown/index.tsx
index 398c939dd0..69b393ed58 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/breadcrumbs/dropdown/index.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/breadcrumbs/dropdown/index.tsx
@@ -44,7 +44,7 @@ const Dropdown = ({
         <button
           type='button'
           className={cn(
-            'rounded-md p-1',
+            'flex size-6 items-center justify-center rounded-md',
             open ? 'bg-state-base-hover' : 'hover:bg-state-base-hover',
           )}
         >
diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/breadcrumbs/index.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/breadcrumbs/index.tsx
index 9c6e325df0..7752b754d2 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/breadcrumbs/index.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/breadcrumbs/index.tsx
@@ -4,6 +4,7 @@ import { useDataSourceStore } from '../../../../store'
 import Bucket from './bucket'
 import BreadcrumbItem from './item'
 import Dropdown from './dropdown'
+import type { OnlineDriveFile } from '@/models/pipeline'
 
 type BreadcrumbsProps = {
   prefix: string[]
@@ -11,6 +12,12 @@ type BreadcrumbsProps = {
   bucket: string
   searchResultsLength: number
   isInPipeline: boolean
+  getOnlineDriveFiles: (params: {
+    prefix?: string[]
+    bucket?: string
+    startAfter?: string
+    fileList?: OnlineDriveFile[]
+  }) => void
 }
 
 const Breadcrumbs = ({
@@ -19,6 +26,7 @@ const Breadcrumbs = ({
   bucket,
   searchResultsLength,
   isInPipeline,
+  getOnlineDriveFiles,
 }: BreadcrumbsProps) => {
   const { t } = useTranslation()
   const dataSourceStore = useDataSourceStore()
@@ -48,14 +56,23 @@ const Breadcrumbs = ({
     setSelectedFileList([])
     setBucket('')
     setPrefix([])
-  }, [dataSourceStore])
+    getOnlineDriveFiles({
+      prefix: [],
+      bucket: '',
+      fileList: [],
+    })
+  }, [dataSourceStore, getOnlineDriveFiles])
 
   const handleClickBucketName = useCallback(() => {
     const { setFileList, setSelectedFileList, setPrefix } = dataSourceStore.getState()
     setFileList([])
     setSelectedFileList([])
     setPrefix([])
-  }, [dataSourceStore])
+    getOnlineDriveFiles({
+      prefix: [],
+      fileList: [],
+    })
+  }, [dataSourceStore, getOnlineDriveFiles])
 
   const handleClickBreadcrumb = useCallback((index: number) => {
     const { setFileList, setSelectedFileList, setPrefix } = dataSourceStore.getState()
@@ -63,12 +80,16 @@ const Breadcrumbs = ({
     setFileList([])
     setSelectedFileList([])
     setPrefix(newPrefix)
-  }, [dataSourceStore, prefix])
+    getOnlineDriveFiles({
+      prefix: newPrefix,
+      fileList: [],
+    })
+  }, [dataSourceStore, getOnlineDriveFiles, prefix])
 
   return (
-    <div className='flex grow items-center overflow-hidden py-1'>
+    <div className='flex grow items-center overflow-hidden'>
       {showSearchResult && (
-        <div className='system-sm-medium text-test-secondary px-[5px] py-1'>
+        <div className='system-sm-medium text-test-secondary px-[5px]'>
           {t('datasetPipeline.onlineDrive.breadcrumbs.searchResult', {
             searchResultsLength,
             folderName: prefix.length > 0 ? prefix[prefix.length - 1] : bucket,
@@ -76,7 +97,7 @@ const Breadcrumbs = ({
         </div>
       )}
       {!showSearchResult && isRoot && (
-        <div className='system-sm-medium text-test-secondary px-[5px] py-1'>
+        <div className='system-sm-medium text-test-secondary px-[5px]'>
           {t('datasetPipeline.onlineDrive.breadcrumbs.allBuckets')}
         </div>
       )}
diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/index.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/index.tsx
index 5cdeb5faaa..174e44ea78 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/index.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/header/index.tsx
@@ -2,6 +2,7 @@ import React from 'react'
 import Breadcrumbs from './breadcrumbs'
 import Input from '@/app/components/base/input'
 import { useTranslation } from 'react-i18next'
+import type { OnlineDriveFile } from '@/models/pipeline'
 
 type HeaderProps = {
   prefix: string[]
@@ -12,6 +13,12 @@ type HeaderProps = {
   handleInputChange: React.ChangeEventHandler<HTMLInputElement>
   handleResetKeywords: () => void
   isInPipeline: boolean
+  getOnlineDriveFiles: (params: {
+    prefix?: string[]
+    bucket?: string
+    startAfter?: string
+    fileList?: OnlineDriveFile[]
+  }) => void
 }
 
 const Header = ({
@@ -23,6 +30,7 @@ const Header = ({
   searchResultsLength,
   handleInputChange,
   handleResetKeywords,
+  getOnlineDriveFiles,
 }: HeaderProps) => {
   const { t } = useTranslation()
 
@@ -34,6 +42,7 @@ const Header = ({
         bucket={bucket}
         searchResultsLength={searchResultsLength}
         isInPipeline={isInPipeline}
+        getOnlineDriveFiles={getOnlineDriveFiles}
       />
       <Input
         value={inputValue}
diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/index.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/index.tsx
index c47509fb63..7db046e13f 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/index.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/index.tsx
@@ -18,6 +18,12 @@ type FileListProps = {
   handleOpenFolder: (file: OnlineDriveFile) => void
   isLoading: boolean
   isTruncated: boolean
+  getOnlineDriveFiles: (params: {
+    prefix?: string[]
+    bucket?: string
+    startAfter?: string
+    fileList?: OnlineDriveFile[]
+  }) => void
 }
 
 const FileList = ({
@@ -34,6 +40,7 @@ const FileList = ({
   isInPipeline,
   isLoading,
   isTruncated,
+  getOnlineDriveFiles,
 }: FileListProps) => {
   const [inputValue, setInputValue] = useState(keywords)
 
@@ -66,6 +73,7 @@ const FileList = ({
         handleInputChange={handleInputChange}
         searchResultsLength={searchResultsLength}
         handleResetKeywords={handleResetKeywords}
+        getOnlineDriveFiles={getOnlineDriveFiles}
       />
       <List
         fileList={fileList}
@@ -77,6 +85,7 @@ const FileList = ({
         isInPipeline={isInPipeline}
         isLoading={isLoading}
         isTruncated={isTruncated}
+        getOnlineDriveFiles={getOnlineDriveFiles}
       />
     </div>
   )
diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/list/index.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/list/index.tsx
index 3e290999b6..22c3f408a5 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/list/index.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/list/index.tsx
@@ -19,6 +19,12 @@ type FileListProps = {
   handleResetKeywords: () => void
   handleSelectFile: (file: OnlineDriveFile) => void
   handleOpenFolder: (file: OnlineDriveFile) => void
+  getOnlineDriveFiles: (params: {
+    prefix?: string[]
+    bucket?: string
+    startAfter?: string
+    fileList?: OnlineDriveFile[]
+  }) => void
 }
 
 const List = ({
@@ -31,6 +37,7 @@ const List = ({
   isInPipeline,
   isLoading,
   isTruncated,
+  getOnlineDriveFiles,
 }: FileListProps) => {
   const anchorRef = useRef<HTMLDivElement>(null)
   const observerRef = useRef<IntersectionObserver>()
@@ -40,15 +47,18 @@ const List = ({
     if (anchorRef.current) {
       observerRef.current = new IntersectionObserver((entries) => {
         const { setStartAfter } = dataSourceStore.getState()
-        if (entries[0].isIntersecting && isTruncated && !isLoading)
+        if (entries[0].isIntersecting && isTruncated && !isLoading) {
           setStartAfter(fileList[fileList.length - 1].key)
+          getOnlineDriveFiles({ startAfter: fileList[fileList.length - 1].key })
+        }
       }, {
         rootMargin: '100px',
       })
       observerRef.current.observe(anchorRef.current)
     }
     return () => observerRef.current?.disconnect()
-  }, [anchorRef, dataSourceStore, isTruncated, isLoading, fileList])
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [anchorRef])
 
   const isAllLoading = isLoading && fileList.length === 0 && keywords.length === 0
   const isPartLoading = isLoading && fileList.length > 0
diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/index.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/index.tsx
index 052120a4b2..f33d03d348 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/index.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/index.tsx
@@ -38,8 +38,17 @@ const OnlineDrive = ({
     ? `/rag/pipelines/${pipelineId}/workflows/published/datasource/nodes/${nodeId}/run`
     : `/rag/pipelines/${pipelineId}/workflows/draft/datasource/nodes/${nodeId}/run`
 
-  const getOnlineDrive = useCallback(async () => {
-    const prefixString = prefix.length > 0 ? `${prefix.join('/')}/` : ''
+  const getOnlineDriveFiles = useCallback(async (params: {
+    prefix?: string[]
+    bucket?: string
+    startAfter?: string
+    fileList?: OnlineDriveFile[]
+  }) => {
+    const _prefix = params.prefix ?? prefix
+    const _bucket = params.bucket ?? bucket
+    const _startAfter = params.startAfter ?? startAfter
+    const _fileList = params.fileList ?? fileList
+    const prefixString = _prefix.length > 0 ? `${_prefix.join('/')}/` : ''
     setIsLoading(true)
     ssePost(
       datasourceNodeRunURL,
@@ -47,8 +56,8 @@ const OnlineDrive = ({
         body: {
           inputs: {
             prefix: prefixString,
-            bucket,
-            start_after: startAfter,
+            bucket: _bucket,
+            start_after: _startAfter,
             max_keys: 30, // Adjust as needed
           },
           datasource_type: DatasourceType.onlineDrive,
@@ -57,8 +66,8 @@ const OnlineDrive = ({
       {
         onDataSourceNodeCompleted: (documentsData: DataSourceNodeCompletedResponse) => {
           const { setFileList, setIsTruncated } = dataSourceStore.getState()
-          const { fileList: newFileList, isTruncated } = convertOnlineDriveData(documentsData.data, prefix)
-          setFileList([...fileList, ...newFileList])
+          const { fileList: newFileList, isTruncated } = convertOnlineDriveData(documentsData.data, _prefix)
+          setFileList([..._fileList, ...newFileList])
           setIsTruncated(isTruncated)
           setIsLoading(false)
         },
@@ -71,12 +80,13 @@ const OnlineDrive = ({
         },
       },
     )
-  }, [prefix, datasourceNodeRunURL, bucket, startAfter, dataSourceStore, fileList])
+  }, [prefix, bucket, startAfter, datasourceNodeRunURL, dataSourceStore, fileList])
 
   useEffect(() => {
-    getOnlineDrive()
+    if (fileList.length > 0) return
+    getOnlineDriveFiles({})
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [bucket, prefix, startAfter])
+  }, [])
 
   const onlineDriveFileList = useMemo(() => {
     if (keywords)
@@ -117,6 +127,7 @@ const OnlineDrive = ({
     setFileList([])
     if (file.type === OnlineDriveFileType.bucket) {
       setBucket(file.displayName)
+      getOnlineDriveFiles({ bucket: file.displayName, fileList: [] })
     }
     else {
       setSelectedFileList([])
@@ -125,8 +136,9 @@ const OnlineDrive = ({
         draft.push(displayName)
       })
       setPrefix(newPrefix)
+      getOnlineDriveFiles({ prefix: newPrefix, fileList: [] })
     }
-  }, [dataSourceStore, prefix])
+  }, [dataSourceStore, getOnlineDriveFiles, prefix])
 
   return (
     <div className='flex flex-col gap-y-2'>
@@ -148,6 +160,7 @@ const OnlineDrive = ({
         isInPipeline={isInPipeline}
         isLoading={isLoading}
         isTruncated={isTruncated}
+        getOnlineDriveFiles={getOnlineDriveFiles}
       />
     </div>
   )
diff --git a/web/app/components/datasets/documents/create-from-pipeline/index.tsx b/web/app/components/datasets/documents/create-from-pipeline/index.tsx
index 539224883c..6b16a5748b 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/index.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/index.tsx
@@ -113,17 +113,6 @@ const CreateFormPipeline = () => {
     return false
   }, [datasource, datasourceType, isShowVectorSpaceFull, fileList.length, allFileLoaded, onlineDocuments.length, websitePages.length, selectedFileList.length])
 
-  const showSelect = useMemo(() => {
-    if (datasourceType === DatasourceType.onlineDocument) {
-      const pagesCount = currentWorkspace?.pages.length ?? 0
-      return pagesCount > 0
-    }
-    if (datasourceType === DatasourceType.onlineDrive) {
-      const isBucketList = onlineDriveFileList.some(file => file.type === 'bucket')
-      return !isBucketList && onlineDriveFileList.length > 0
-    }
-  }, [currentWorkspace?.pages.length, datasourceType, onlineDriveFileList])
-
   const supportedFileTypes = useMemo(() => {
     if (!supportFileTypesRes) return []
     return Array.from(new Set(supportFileTypesRes.allowed_extensions.map(item => item.toLowerCase())))
@@ -134,6 +123,23 @@ const CreateFormPipeline = () => {
     batch_count_limit: 5,
   }, [fileUploadConfigResponse])
 
+  const showSelect = useMemo(() => {
+    if (datasourceType === DatasourceType.onlineDocument) {
+      const pagesCount = currentWorkspace?.pages.length ?? 0
+      return pagesCount > 0
+    }
+    if (datasourceType === DatasourceType.onlineDrive) {
+      const isBucketList = onlineDriveFileList.some(file => file.type === 'bucket')
+      return !isBucketList && onlineDriveFileList.filter((item) => {
+        if (item.type === 'bucket') return false
+        if (item.type === 'folder') return true
+        if (item.type === 'file')
+          return supportedFileTypes.includes(getFileExtension(item.key))
+        return false
+      }).length > 0
+    }
+  }, [currentWorkspace?.pages.length, datasourceType, onlineDriveFileList, supportedFileTypes])
+
   const totalOptions = useMemo(() => {
     if (datasourceType === DatasourceType.onlineDocument)
       return currentWorkspace?.pages.length
diff --git a/web/app/components/datasets/list/datasets.tsx b/web/app/components/datasets/list/datasets.tsx
index 70d772466a..0e090e72ee 100644
--- a/web/app/components/datasets/list/datasets.tsx
+++ b/web/app/components/datasets/list/datasets.tsx
@@ -52,7 +52,8 @@ const Datasets = ({
       observerRef.current.observe(anchorRef.current)
     }
     return () => observerRef.current?.disconnect()
-  }, [anchorRef, datasetList, hasNextPage, fetchNextPage, isFetching])
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [anchorRef])
 
   return (
     <>
diff --git a/web/app/components/rag-pipeline/components/panel/test-run/hooks.ts b/web/app/components/rag-pipeline/components/panel/test-run/hooks.ts
index 75d3220432..59de068304 100644
--- a/web/app/components/rag-pipeline/components/panel/test-run/hooks.ts
+++ b/web/app/components/rag-pipeline/components/panel/test-run/hooks.ts
@@ -5,7 +5,6 @@ import { useNodes } from 'reactflow'
 import { BlockEnum } from '@/app/components/workflow/types'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
 import { useCallback, useMemo, useState } from 'react'
-import type { OnlineDriveFile } from '@/models/pipeline'
 
 export const useTestRunSteps = () => {
   const { t } = useTranslation()
@@ -57,27 +56,3 @@ export const useDatasourceOptions = () => {
 
   return options
 }
-
-export const useOnlineDrive = () => {
-  const [prefix, setPrefix] = useState<string[]>([])
-  const [keywords, setKeywords] = useState('')
-  const [bucket, setBucket] = useState('')
-  const [startAfter, setStartAfter] = useState('')
-  const [selectedFileList, setSelectedFileList] = useState<string[]>([])
-  const [fileList, setFileList] = useState<OnlineDriveFile[]>([])
-
-  return {
-    prefix,
-    setPrefix,
-    keywords,
-    setKeywords,
-    bucket,
-    setBucket,
-    startAfter,
-    setStartAfter,
-    selectedFileList,
-    setSelectedFileList,
-    fileList,
-    setFileList,
-  }
-}