diff --git a/ui/src/components/pages/Pipeline/partials/Graph/index.test.tsx b/ui/src/components/pages/Pipeline/partials/Graph/index.test.tsx
index 0d979124c3..7af2f6f994 100644
--- a/ui/src/components/pages/Pipeline/partials/Graph/index.test.tsx
+++ b/ui/src/components/pages/Pipeline/partials/Graph/index.test.tsx
@@ -2,9 +2,61 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { Position } from "reactflow";
 import Graph from "./index";
 
-global.ResizeObserver = require("resize-observer-polyfill");
+class ResizeObserver {
+  callback: globalThis.ResizeObserverCallback;
+  constructor(callback: globalThis.ResizeObserverCallback) {
+    this.callback = callback;
+  }
+  observe(target: Element) {
+    this.callback([{ target } as globalThis.ResizeObserverEntry], this);
+  }
+  unobserve() {
+    return;
+  }
+  disconnect() {
+    return;
+  }
+}
+
+class DOMMatrixReadOnly {
+  m22: number;
+  constructor(transform: string) {
+    const scale = transform?.match(/scale\(([1-9.])\)/)?.[1];
+    this.m22 = scale !== undefined ? +scale : 1;
+  }
+}
+
+let init = false;
+
+export const mockReactFlow = () => {
+  if (init) return;
+  init = true;
+  global.ResizeObserver = ResizeObserver;
+  global.DOMMatrixReadOnly = DOMMatrixReadOnly;
+  Object.defineProperties(global.HTMLElement.prototype, {
+    offsetHeight: {
+      get() {
+        return parseFloat(this.style.height) || 1;
+      },
+    },
+    offsetWidth: {
+      get() {
+        return parseFloat(this.style.width) || 1;
+      },
+    },
+  });
+  (global.SVGElement as any).prototype.getBBox = () => ({
+    x: 0,
+    y: 0,
+    width: 0,
+    height: 0,
+  });
+};
 
 describe("Graph screen test", () => {
+  beforeEach(() => {
+    mockReactFlow();
+  });
   const data = {
     edges: [
       {
@@ -157,6 +209,7 @@ describe("Graph screen test", () => {
     await waitFor(() => {
       expect(screen.getByTestId("graph")).toBeVisible();
       fireEvent.click(screen.getByTestId("rf__node-in"));
+      fireEvent.click(screen.getByTestId("rf__edge-in-cat"));
       fireEvent.click(container.getElementsByClassName("react-flow__pane")[0]);
     });
     await waitFor(() => expect(screen.getByTestId("card")).toBeVisible());
diff --git a/ui/src/components/pages/Pipeline/partials/Graph/partials/CustomEdge/index.test.tsx b/ui/src/components/pages/Pipeline/partials/Graph/partials/CustomEdge/index.test.tsx
index a584a5c552..8eecb20672 100644
--- a/ui/src/components/pages/Pipeline/partials/Graph/partials/CustomEdge/index.test.tsx
+++ b/ui/src/components/pages/Pipeline/partials/Graph/partials/CustomEdge/index.test.tsx
@@ -4,7 +4,7 @@ import { ReactFlowProvider } from "reactflow";
 import CustomEdge from "./index";
 
 describe("Graph screen test", () => {
-  it("Straight edge not full", async () => {
+  it("Fwd edge not full", async () => {
     render(
       <ReactFlowProvider>
         <CustomEdge
@@ -20,6 +20,8 @@ describe("Graph screen test", () => {
             edgeWatermark: {
               isWaterMarkEnabled: true,
             },
+            fwdEdge: true,
+            fromNodeOutDegree: 1,
           }}
           source={"first"}
           target={"second"}
@@ -30,7 +32,7 @@ describe("Graph screen test", () => {
       expect(screen.getByTestId(`first-second`)).toBeInTheDocument()
     );
   });
-  it("Straight edge but full with delay in mo", async () => {
+  it("Fwd edge but full with delay in mo", async () => {
     render(
       <ReactFlowProvider>
         <CustomEdge
@@ -48,6 +50,7 @@ describe("Graph screen test", () => {
               isWaterMarkEnabled: true,
               watermarks: [Date.now() - 2678400000],
             },
+            fwdEdge: true,
           }}
           source={"first"}
           target={"second"}
@@ -76,6 +79,7 @@ describe("Graph screen test", () => {
               isWaterMarkEnabled: true,
               watermarks: [Date.now() - 1],
             },
+            fwdEdge: true,
           }}
           source={"first"}
           target={"second"}
@@ -95,6 +99,7 @@ describe("Graph screen test", () => {
               isWaterMarkEnabled: true,
               watermarks: [Date.now() - 1000],
             },
+            fwdEdge: true,
           }}
           source={"first"}
           target={"third"}
@@ -114,6 +119,7 @@ describe("Graph screen test", () => {
               isWaterMarkEnabled: true,
               watermarks: [Date.now() - 60000],
             },
+            fwdEdge: true,
           }}
           source={"first"}
           target={"fourth"}
@@ -133,6 +139,7 @@ describe("Graph screen test", () => {
               isWaterMarkEnabled: true,
               watermarks: [Date.now() - 3600000],
             },
+            fwdEdge: true,
           }}
           source={"first"}
           target={"fifth"}
@@ -152,6 +159,7 @@ describe("Graph screen test", () => {
               isWaterMarkEnabled: true,
               watermarks: [Date.now() - 86400000],
             },
+            fwdEdge: true,
           }}
           source={"first"}
           target={"sixth"}
@@ -174,4 +182,114 @@ describe("Graph screen test", () => {
       expect(screen.getByTestId(`first-sixth`)).toBeInTheDocument()
     );
   });
+  it("Forward edges with out-degree > 1", async () => {
+    render(
+      <ReactFlowProvider>
+        <CustomEdge
+          id={"first-second"}
+          sourceX={240}
+          sourceY={36}
+          targetX={334}
+          targetY={40}
+          sourcePosition={Position.Right}
+          targetPosition={Position.Left}
+          data={{
+            isFull: "true",
+            backpressureLabel: "0",
+            edgeWatermark: {
+              isWaterMarkEnabled: true,
+              watermarks: [Date.now() - 2678400000],
+            },
+            fwdEdge: true,
+            fromNodeOutDegree: 2,
+          }}
+          source={"first"}
+          target={"second"}
+        />
+        <CustomEdge
+          id={"first-third"}
+          sourceX={240}
+          sourceY={36}
+          targetX={334}
+          targetY={32}
+          sourcePosition={Position.Right}
+          targetPosition={Position.Left}
+          data={{
+            isFull: "true",
+            backpressureLabel: "0",
+            edgeWatermark: {
+              isWaterMarkEnabled: true,
+              watermarks: [Date.now() - 2678400000],
+            },
+            fwdEdge: true,
+            fromNodeOutDegree: 2,
+          }}
+          source={"first"}
+          target={"third"}
+        />
+      </ReactFlowProvider>
+    );
+    await waitFor(() => {
+      expect(screen.getByTestId(`first-second`)).toBeInTheDocument();
+      expect(screen.getByTestId(`first-third`)).toBeInTheDocument();
+    });
+  });
+  it("Back edges", async () => {
+    render(
+      <ReactFlowProvider>
+        <CustomEdge
+          id={"first-second"}
+          sourceX={334}
+          sourceY={36}
+          targetX={240}
+          targetY={36}
+          sourcePosition={Position.Right}
+          targetPosition={Position.Left}
+          data={{
+            isFull: "true",
+            backpressureLabel: "0",
+            edgeWatermark: {
+              isWaterMarkEnabled: true,
+              watermarks: [Date.now() - 2678400000],
+            },
+            backEdge: true,
+          }}
+          source={"first"}
+          target={"second"}
+        />
+      </ReactFlowProvider>
+    );
+    await waitFor(() =>
+      expect(screen.getByTestId(`first-second`)).toBeInTheDocument()
+    );
+  });
+  it("Self edges", async () => {
+    render(
+      <ReactFlowProvider>
+        <CustomEdge
+          id={"first-first"}
+          sourceX={334}
+          sourceY={36}
+          targetX={300}
+          targetY={36}
+          sourcePosition={Position.Top}
+          targetPosition={Position.Top}
+          data={{
+            isFull: "true",
+            backpressureLabel: "0",
+            edgeWatermark: {
+              isWaterMarkEnabled: true,
+              watermarks: [Date.now() - 2678400000],
+            },
+            selfEdge: true,
+          }}
+          source={"first"}
+          target={"second"}
+        />
+      </ReactFlowProvider>
+    );
+    await waitFor(() =>
+      expect(screen.getByTestId(`first-first`)).toBeInTheDocument()
+    );
+  });
 });
diff --git a/ui/src/components/pages/Pipeline/partials/Graph/partials/CustomEdge/index.tsx b/ui/src/components/pages/Pipeline/partials/Graph/partials/CustomEdge/index.tsx
index 99b7366550..c2b3f5e715 100644
--- a/ui/src/components/pages/Pipeline/partials/Graph/partials/CustomEdge/index.tsx
+++ b/ui/src/components/pages/Pipeline/partials/Graph/partials/CustomEdge/index.tsx
@@ -40,9 +40,9 @@ const CustomEdge: FC<EdgeProps> = ({
         }
       } else {
         if (sourceY < targetY) {
-          labelY = targetY - 47;
+          labelY = sourceY + 1;
         } else {
-          labelY = targetY + 36;
+          labelY = sourceY - 13;
         }
       }
     }