diff --git a/web/app/components/workflow/block-icon.tsx b/web/app/components/workflow/block-icon.tsx
index 4ec78a49bc..234e20c7ff 100644
--- a/web/app/components/workflow/block-icon.tsx
+++ b/web/app/components/workflow/block-icon.tsx
@@ -109,7 +109,7 @@ const BlockIcon: FC<BlockIconProps> = ({
         )
       }
       {
-        type === BlockEnum.Tool && toolIcon && (
+        (type === BlockEnum.Tool || type === BlockEnum.DataSource) && toolIcon && (
           <>
             {
               typeof toolIcon === 'string'