apache · TomasCorreia2003 · May 29, 2026
@@ -92,6 +92,7 @@
     "downloadImage": "Download graph image",
     "downloadImageError": "Failed to download graph image.",
     "downloadImageErrorTitle": "Download Failed",
+    "manualLayout": "Rearrange",
     "otherDagRuns": "+Other Dag Runs",
     "taskCount_one": "{{count}} Task",
     "taskCount_other": "{{count}} Tasks",

@@ -0,0 +1,138 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { render } from "@testing-library/react";
+import { Position } from "@xyflow/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import CustomEdge from "./Edge";
+import type { EdgeData } from "./reactflowUtils";
+
+const chakraMocks = vi.hoisted(() => ({
+  useToken: vi.fn(() => ["#111111", "#222222", "#333333", "#444444", "#555555", "#666666"]),
+}));
+const xyFlowMocks = vi.hoisted(() => ({
+  BaseEdge: vi.fn(() => null),
+  getSmoothStepPath: vi.fn(() => ["M 0 0"]),
+  useNodesData: vi.fn(() => [{ data: { isSelected: false } }, { data: { isSelected: false } }]),
+  useStore: vi.fn(),
+}));
+
+vi.mock("@chakra-ui/react", async () => {
+  const actual = await vi.importActual("@chakra-ui/react");
+
+  return { ...actual, useToken: chakraMocks.useToken };
+});
+
+vi.mock("@xyflow/react", async () => {
+  const actual = await vi.importActual("@xyflow/react");
+
+  return {
+    ...actual,
+    BaseEdge: xyFlowMocks.BaseEdge,
+    getSmoothStepPath: xyFlowMocks.getSmoothStepPath,
+    useNodesData: xyFlowMocks.useNodesData,
+    useStore: xyFlowMocks.useStore,
+  };
+});
+
+type MockNodeLookupEntry = {
+  internals: { userNode: { dragging?: boolean } };
+};
+
+type MockStoreState = {
+  nodeLookup: Map<string, MockNodeLookupEntry>;
+};
+
+let nodeLookup = new Map<string, MockNodeLookupEntry>();
+
+const buildEdgeProps = ({
+  source,
+  target,
+}: {
+  source: string;
+  target: string;
+}): {
+  data: EdgeData;
+  id: string;
+  source: string;
+  sourcePosition: Position.Right;
+  sourceX: number;
+  sourceY: number;
+  target: string;
+  targetPosition: Position.Left;
+  targetX: number;
+  targetY: number;
+  type: "custom";
+} => ({
+  data: {
+    isManualLayout: true,
+    rest: {
+      id: `edge-${source}-${target}`,
+      sources: [source],
+      targets: [target],
+    },
+  },
+  id: `edge-${source}-${target}`,
+  source,
+  sourcePosition: Position.Right,
+  sourceX: 0,
+  sourceY: 0,
+  target,
+  targetPosition: Position.Left,
+  targetX: 100,
+  targetY: 100,
+  type: "custom",
+});
+
+const getLastBaseEdgeStroke = () => {
+  const baseEdgeCalls = vi.mocked(xyFlowMocks.BaseEdge).mock.calls as unknown as Array<
+    [{ style?: { stroke?: string | undefined } | undefined }]
+  >;
+
+  return baseEdgeCalls.at(-1)?.[0]?.style?.stroke;
+};
+
+describe("CustomEdge", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    nodeLookup = new Map([
+      ["task_1", { internals: { userNode: { dragging: true } } }],
+      ["task_2", { internals: { userNode: { dragging: false } } }],
+      ["task_3", { internals: { userNode: { dragging: false } } }],
+      ["task_4", { internals: { userNode: { dragging: false } } }],
+    ]);
+
+    vi.mocked(xyFlowMocks.useStore).mockImplementation((selector: (state: MockStoreState) => boolean) =>
+      selector({ nodeLookup }),
+    );
+  });
+
+  it("uses drag preview colors only for edges connected to dragged nodes", () => {
+    render(<CustomEdge {...buildEdgeProps({ source: "task_1", target: "task_2" })} />);
+
+    const connectedStroke = getLastBaseEdgeStroke();
+
+    render(<CustomEdge {...buildEdgeProps({ source: "task_3", target: "task_4" })} />);
+
+    const unconnectedStroke = getLastBaseEdgeStroke();
+
+    expect(connectedStroke).toBe("#444444");
+    expect(unconnectedStroke).toBe("#111111");
+  });
+});
@@ -19,38 +19,95 @@
 import { Text, useToken } from "@chakra-ui/react";
 import { Group } from "@visx/group";
 import { LinePath } from "@visx/shape";
-import type { Edge as EdgeType } from "@xyflow/react";
-import { useNodesData } from "@xyflow/react";
+import { BaseEdge, getSmoothStepPath, useNodesData, useStore } from "@xyflow/react";
+import type { Edge as EdgeType, EdgeProps } from "@xyflow/react";
 import type { ElkPoint } from "elkjs";
 
 import { opacityStyle } from "./graphTypes";
 import type { EdgeData } from "./reactflowUtils";
 
-type Props = EdgeType<EdgeData>;
+type Props = EdgeProps<EdgeType<EdgeData>>;
 
-const CustomEdge = ({ data, source, target }: Props) => {
-  const [strokeColor, blueColor, dataEdgeColor] = useToken("colors", [
-    "border.inverted",
-    "blue.500",
-    "purple.500",
-  ]);
+const CustomEdge = ({
+  data,
+  id,
+  markerEnd,
+  markerStart,
+  source,
+  sourcePosition,
+  sourceX,
+  sourceY,
+  target,
+  targetPosition,
+  targetX,
+  targetY,
+}: Props) => {
+  const [
+    strokeColor,
+    blueColor,
+    dataEdgeColor,
+    draggingStrokeColor,
+    draggingBlueColor,
+    draggingDataEdgeColor,
+  ] = useToken("colors", ["border.inverted", "blue.500", "purple.500", "border", "blue.400", "purple.400"]);
 
   // Read isSelected directly from the node store so that selection changes
   // don't require the parent to rebuild and pass down a new edges array.
   // useNodesData subscribes to data changes for these specific node IDs only.
   const nodesData = useNodesData([source, target]);
   const isSelected = nodesData.some((node) => Boolean(node.data.isSelected));
+  const isConnectedNodeDragging = useStore((state) => {
+    const sourceDragging = state.nodeLookup.get(source)?.internals.userNode.dragging ?? false;
+    const targetDragging = state.nodeLookup.get(target)?.internals.userNode.dragging ?? false;
+
+    return sourceDragging || targetDragging;
+  });
 
   if (data === undefined) {
     return undefined;
   }
-  const { rest } = data;
+  const { isManualLayout = false, rest } = data;
+  const isDragPreview = isManualLayout && isConnectedNodeDragging;
+  const selectedEdgeStrokeColor = rest.edgeType === "data" ? dataEdgeColor : blueColor;
+  const selectedDraggingEdgeStrokeColor =
+    rest.edgeType === "data" ? draggingDataEdgeColor : draggingBlueColor;
+  let edgeStrokeColor = (isSelected ? selectedEdgeStrokeColor : strokeColor) ?? "currentColor";
+
+  if (isDragPreview) {
+    edgeStrokeColor = (isSelected ? selectedDraggingEdgeStrokeColor : draggingStrokeColor) ?? "currentColor";
+  }
 
-  const edgeStrokeColor = isSelected ? (rest.edgeType === "data" ? dataEdgeColor : blueColor) : strokeColor;
+  if (isManualLayout) {
+    const [path] = getSmoothStepPath({
+      sourcePosition,
+      sourceX,
+      sourceY,
+      targetPosition,
+      targetX,
+      targetY,
+    });
+
+    return (
+      <g {...opacityStyle(rest.isFiltered)} pointerEvents="none">
+        <BaseEdge
+          id={id}
+          interactionWidth={0}
+          markerEnd={markerEnd}
+          markerStart={markerStart}
+          path={path}
+          style={{
+            stroke: edgeStrokeColor,
+            strokeDasharray: rest.isSetupTeardown ? "10,5" : undefined,
+            strokeWidth: isSelected ? 3 : 1,
+          }}
+        />
+      </g>
+    );
+  }
 
   return (
-    <g {...opacityStyle(rest.isFiltered)}>
-      {rest.labels?.map(({ height, id, text, width, x, y }) => {
+    <g {...opacityStyle(rest.isFiltered)} pointerEvents="none">
+      {rest.labels?.map(({ height, id: labelId, text, width, x, y }) => {
         if (y === undefined || x === undefined) {
           return undefined;
         }
@@ -59,7 +116,7 @@ const CustomEdge = ({ data, source, target }: Props) => {
           <Group
             // Add a tiny bit of height so letters aren't cut off
             height={(height ?? 0) + 2}
-            key={id}
+            key={labelId}
             left={x}
             top={y}
             width={width}

@@ -122,6 +122,7 @@ type Edge = {
 } & ElkExtendedEdge;
 
 export type EdgeData = {
+  isManualLayout?: boolean;
   rest: {
     edgeType?: "data" | "scheduling";
     isFiltered?: boolean;