diff --git a/ui/src/components/workflow/node/BranchNode.tsx b/ui/src/components/workflow/node/BranchNode.tsx
index 61eb0c9e..4a68f315 100644
--- a/ui/src/components/workflow/node/BranchNode.tsx
+++ b/ui/src/components/workflow/node/BranchNode.tsx
@@ -1,6 +1,6 @@
 import { memo } from "react";
 import { useTranslation } from "react-i18next";
-import { Button } from "antd";
+import { Button, theme } from "antd";
 
 import { type WorkflowNode } from "@/domain/workflow";
 import { useZustandShallowSelector } from "@/hooks";
@@ -10,6 +10,8 @@ import AddNode from "./AddNode";
 import WorkflowElement from "../WorkflowElement";
 import { type SharedNodeProps } from "./_SharedNode";
 
+const { useToken } = theme;
+
 export type BrandNodeProps = SharedNodeProps;
 
 const BranchNode = ({ node, disabled }: BrandNodeProps) => {
@@ -17,6 +19,8 @@ const BranchNode = ({ node, disabled }: BrandNodeProps) => {
 
   const { addBranch } = useWorkflowStore(useZustandShallowSelector(["addBranch"]));
 
+  const token = useToken();
+
   const renderBranch = (node: WorkflowNode, branchNodeId?: string, branchIndex?: number) => {
     const elements: JSX.Element[] = [];
 
@@ -31,7 +35,12 @@ const BranchNode = ({ node, disabled }: BrandNodeProps) => {
 
   return (
     <>
-      <div className="relative flex gap-x-16 before:absolute before:inset-x-[128px] before:top-0 before:h-[2px] before:bg-stone-200 before:content-[''] after:absolute after:inset-x-[128px] after:bottom-0 after:h-[2px] after:bg-stone-200 after:content-['']">
+      <div
+        className="relative flex gap-x-16 before:absolute before:inset-x-[128px] before:top-0 before:h-[2px] before:bg-stone-200 before:content-[''] after:absolute after:inset-x-[128px] after:bottom-0 after:h-[2px] after:bg-stone-200 after:content-['']"
+        style={{
+          backgroundColor: token.token.colorBgContainer,
+        }}
+      >
         <Button
           className="absolute left-1/2 z-[1] -translate-x-1/2 -translate-y-1/2 text-xs"
           disabled={disabled}
@@ -50,6 +59,38 @@ const BranchNode = ({ node, disabled }: BrandNodeProps) => {
             key={branch.id}
             className="relative flex flex-col items-center before:absolute  before:left-1/2 before:top-0 before:h-full before:w-[2px] before:-translate-x-1/2 before:bg-stone-200 before:content-['']"
           >
+            {index == 0 && (
+              <>
+                <div
+                  className="absolute -left-px -top-1 h-2 w-1/2"
+                  style={{
+                    backgroundColor: token.token.colorBgContainer,
+                  }}
+                ></div>
+                <div
+                  className="absolute -bottom-1 -left-px z-50 h-2 w-1/2"
+                  style={{
+                    backgroundColor: token.token.colorBgContainer,
+                  }}
+                ></div>
+              </>
+            )}
+            {node.branches && index == node.branches.length - 1 && (
+              <>
+                <div
+                  className="absolute -right-px -top-1 h-2 w-1/2"
+                  style={{
+                    backgroundColor: token.token.colorBgContainer,
+                  }}
+                ></div>
+                <div
+                  className="absolute -bottom-1 -right-px z-50 h-2 w-1/2"
+                  style={{
+                    backgroundColor: token.token.colorBgContainer,
+                  }}
+                ></div>
+              </>
+            )}
             <div className="relative flex flex-col items-center">{renderBranch(branch, node.id, index)}</div>
           </div>
         ))}