diff --git a/web/app/components/rag-pipeline/components/panel/test-run/data-source-options/index.tsx b/web/app/components/rag-pipeline/components/panel/test-run/data-source-options/index.tsx
index cacb0e9d6b..9c00aa1d7d 100644
--- a/web/app/components/rag-pipeline/components/panel/test-run/data-source-options/index.tsx
+++ b/web/app/components/rag-pipeline/components/panel/test-run/data-source-options/index.tsx
@@ -1,11 +1,6 @@
 import { useCallback, useEffect } from 'react'
 import { useDatasourceOptions } from '../hooks'
 import OptionCard from './option-card'
-import { File, Watercrawl } from '@/app/components/base/icons/src/public/knowledge'
-import { Notion } from '@/app/components/base/icons/src/public/common'
-import { Jina } from '@/app/components/base/icons/src/public/llm'
-import { DataSourceType } from '@/models/datasets'
-import { DataSourceProvider } from '@/models/common'
 import type { Datasource } from '../types'
 
 type DataSourceOptionsProps = {
@@ -13,17 +8,6 @@ type DataSourceOptionsProps = {
   onSelect: (option: Datasource) => void
 }
 
-// TODO: replace all icons with tool icon_url
-const DATA_SOURCE_ICONS = {
-  [DataSourceType.FILE]: File as React.FC<React.SVGProps<SVGSVGElement>>,
-  [DataSourceType.NOTION]: Notion as React.FC<React.SVGProps<SVGSVGElement>>,
-  [DataSourceProvider.fireCrawl]: '🔥',
-  [DataSourceProvider.jinaReader]: Jina as React.FC<React.SVGProps<SVGSVGElement>>,
-  [DataSourceProvider.waterCrawl]: Watercrawl as React.FC<React.SVGProps<SVGSVGElement>>,
-}
-
-type KeyType = keyof typeof DATA_SOURCE_ICONS
-
 const DataSourceOptions = ({
   datasourceNodeId,
   onSelect,
@@ -49,8 +33,8 @@ const DataSourceOptions = ({
         <OptionCard
           key={option.value}
           label={option.label}
+          nodeData={option.data}
           selected={datasourceNodeId === option.value}
-          Icon={DATA_SOURCE_ICONS[option.type as KeyType]}
           onClick={handelSelect.bind(null, option.value)}
         />
       ))}
diff --git a/web/app/components/rag-pipeline/components/panel/test-run/data-source-options/option-card.tsx b/web/app/components/rag-pipeline/components/panel/test-run/data-source-options/option-card.tsx
index efc4a42905..0e1dc0b3ff 100644
--- a/web/app/components/rag-pipeline/components/panel/test-run/data-source-options/option-card.tsx
+++ b/web/app/components/rag-pipeline/components/panel/test-run/data-source-options/option-card.tsx
@@ -1,19 +1,25 @@
 import React from 'react'
 import cn from '@/utils/classnames'
+import BlockIcon from '@/app/components/workflow/block-icon'
+import { BlockEnum } from '@/app/components/workflow/types'
+import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
+import { useToolIcon } from '@/app/components/workflow/hooks'
 
 type OptionCardProps = {
   label: string
-  Icon: React.FC<React.SVGProps<SVGSVGElement>> | string
   selected: boolean
+  nodeData: DataSourceNodeType
   onClick?: () => void
 }
 
 const OptionCard = ({
   label,
-  Icon,
   selected,
+  nodeData,
   onClick,
 }: OptionCardProps) => {
+  const toolIcon = useToolIcon(nodeData)
+
   return (
     <div
       className={cn(
@@ -25,11 +31,10 @@ const OptionCard = ({
       onClick={onClick}
     >
       <div className='flex size-7 items-center justify-center rounded-lg border-[0.5px] border-components-panel-border bg-background-default-dodge p-1'>
-        {
-          typeof Icon === 'string'
-            ? <div className='text-[18px] leading-[18px]'>{Icon}</div>
-            : <Icon className='size-5' />
-        }
+        <BlockIcon
+          type={BlockEnum.DataSource}
+          toolIcon={toolIcon}
+        />
       </div>
       <div className={cn('system-sm-medium text-text-secondary', selected && 'text-primary')}>
         {label}
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 986561b8fd..76f7cdd2ac 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
@@ -32,7 +32,7 @@ export const useDatasourceOptions = () => {
         type: node.data.provider_type as DatasourceType,
         variables: node.data.variables || [],
         description: node.data.desc || '',
-        docTitle: '', // todo: Add docTitle and docLink if needed, or remove these properties if not used
+        docTitle: '', // todo: Add docTitle and docLink
         docLink: '',
       }
     })
@@ -41,12 +41,11 @@ export const useDatasourceOptions = () => {
   const options = useMemo(() => {
     const options: DataSourceOption[] = []
     datasourceNodes.forEach((node) => {
-      const type = node.data.provider_type as DatasourceType
       const label = node.data.title
       options.push({
         label,
         value: node.id,
-        type,
+        data: node.data,
       })
     })
     return options
diff --git a/web/app/components/rag-pipeline/components/panel/test-run/types.ts b/web/app/components/rag-pipeline/components/panel/test-run/types.ts
index 61c987a338..9d9bc33df1 100644
--- a/web/app/components/rag-pipeline/components/panel/test-run/types.ts
+++ b/web/app/components/rag-pipeline/components/panel/test-run/types.ts
@@ -1,3 +1,4 @@
+import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
 import type { DatasourceType, RAGPipelineVariables } from '@/models/pipeline'
 
 export enum TestRunStep {
@@ -8,7 +9,7 @@ export enum TestRunStep {
 export type DataSourceOption = {
   label: string
   value: string
-  type: DatasourceType
+  data: DataSourceNodeType
 }
 
 export type Datasource = {
diff --git a/web/models/pipeline.ts b/web/models/pipeline.ts
index 592aa16566..beccc8348e 100644
--- a/web/models/pipeline.ts
+++ b/web/models/pipeline.ts
@@ -7,9 +7,9 @@ import type { Viewport } from 'reactflow'
 import type { TransferMethod } from '@/types/app'
 
 export enum DatasourceType {
-  localFile = 'local-file',
-  onlineDocument = 'online-document',
-  websiteCrawl = 'website-crawl',
+  localFile = 'local_file',
+  onlineDocument = 'online_document',
+  websiteCrawl = 'website_crawl',
 }
 
 export type PipelineTemplateListParams = {