diff --git a/web/app/components/workflow/index.tsx b/web/app/components/workflow/index.tsx
index e4d9ae21a0..56d4b272e8 100644
--- a/web/app/components/workflow/index.tsx
+++ b/web/app/components/workflow/index.tsx
@@ -168,7 +168,7 @@ const WorkflowWrap = memo(() => {
   }, [data])
   const edgesData = useMemo(() => {
     if (data)
-      return initialEdges(data.graph.edges)
+      return initialEdges(data.graph.edges, data.graph.nodes)
 
     return []
   }, [data])
diff --git a/web/app/components/workflow/utils.ts b/web/app/components/workflow/utils.ts
index 4e16619cd3..12ee876009 100644
--- a/web/app/components/workflow/utils.ts
+++ b/web/app/components/workflow/utils.ts
@@ -42,10 +42,21 @@ export const initialNodes = (nodes: Node[], edges: Edge[]) => {
   })
 }
 
-export const initialEdges = (edges: Edge[]) => {
+export const initialEdges = (edges: Edge[], nodes: Node[]) => {
+  const nodesMap = nodes.reduce((acc, node) => {
+    acc[node.id] = node
+
+    return acc
+  }, {} as Record<string, Node>)
   return edges.map((edge) => {
     edge.type = 'custom'
 
+    if (!edge.data?.sourceType)
+      edge.data = { ...edge.data, sourceType: nodesMap[edge.source].data.type! } as any
+
+    if (!edge.data?.targetType)
+      edge.data = { ...edge.data, targetType: nodesMap[edge.target].data.type! } as any
+
     return edge
   })
 }