"use client";

import {
  Background,
  Handle,
  Node,
  NodeProps,
  Position,
  ReactFlow,
  useReactFlow,
} from "@xyflow/react";

import { DataEdge } from "@/components/data-edge";
import { useCallback } from "react";
import { BaseNode } from "../base-node";
import { Button } from "@/components/ui/button";
import { Minus, Plus } from "lucide-react";

const defaultNodes = [
  {
    id: "1",
    position: { x: 100, y: 100 },
    type: "counterNode",
    data: { value: 10 },
  },
  {
    id: "2",
    position: { x: 300, y: 300 },
    data: { label: "Output" },
  },
];

const nodeTypes = {
  counterNode: CounterNode,
};

const defaultEdges = [
  {
    id: "1->2",
    source: "1",
    target: "2",
    type: "dataEdge",
    data: { key: "value" },
  } satisfies DataEdge<CounterNode>,
];

const edgeTypes = {
  dataEdge: DataEdge,
};

export default function DataEdgeDemo() {
  return (
    <div className="h-full w-full">
      <ReactFlow
        defaultNodes={defaultNodes}
        nodeTypes={nodeTypes}
        defaultEdges={defaultEdges}
        edgeTypes={edgeTypes}
        fitView
      >
        <Background />
      </ReactFlow>
    </div>
  );
}

type CounterNode = Node<{ value: number }>;

function CounterNode({ id, data }: NodeProps<CounterNode>) {
  const { updateNodeData } = useReactFlow();
  const handleIncr = useCallback(() => {
    updateNodeData(id, ({ data }) => {
      if ("value" in data && typeof data.value === "number") {
        return { ...data, value: data.value + 1 };
      }

      return data;
    });
  }, [id, updateNodeData]);
  const handleDecr = useCallback(() => {
    updateNodeData(id, ({ data }) => {
      if ("value" in data && typeof data.value === "number") {
        return { ...data, value: data.value - 1 };
      }

      return data;
    });
  }, [id, updateNodeData]);

  return (
    <BaseNode className="flex items-center gap-4">
      <Button onClick={handleDecr} className="nopan nodrag size-6 p-1">
        <Minus />
      </Button>
      <pre>{String(data.value).padStart(2, " ")}</pre>
      <Button onClick={handleIncr} className="nopan nodrag size-6 p-1">
        <Plus />
      </Button>
      <Handle type="source" position={Position.Bottom} />
    </BaseNode>
  );
}


"use client";

import {
  BaseEdge,
  Edge,
  EdgeLabelRenderer,
  EdgeProps,
  getBezierPath,
  getSmoothStepPath,
  getStraightPath,
  Node,
  Position,
  useStore,
} from "@xyflow/react";
import { useMemo } from "react";

export type DataEdge<T extends Node = Node> = Edge<{
  /**
   * The key to lookup in the source node's `data` object. For additional safety,
   * you can parameterize the `DataEdge` over the type of one of your nodes to
   * constrain the possible values of this key.
   *
   * If no key is provided this edge behaves identically to React Flow's default
   * edge component.
   */
  key?: keyof T["data"];
  /**
   * Which of React Flow's path algorithms to use. Each value corresponds to one
   * of React Flow's built-in edge types.
   *
   * If not provided, this defaults to `"bezier"`.
   */
  path?: "bezier" | "smoothstep" | "step" | "straight";
}>;

export function DataEdge({
  data = { path: "bezier" },
  id,
  markerEnd,
  source,
  sourcePosition,
  sourceX,
  sourceY,
  style,
  targetPosition,
  targetX,
  targetY,
}: EdgeProps<DataEdge>) {
  const nodeData = useStore((state) => state.nodeLookup.get(source)?.data);
  const [edgePath, labelX, labelY] = getPath({
    type: data.path ?? "bezier",
    sourceX,
    sourceY,
    sourcePosition,
    targetX,
    targetY,
    targetPosition,
  });

  const label = useMemo(() => {
    if (data.key && nodeData) {
      const value = nodeData[data.key];

      switch (typeof value) {
        case "string":
        case "number":
          return value;

        case "object":
          return JSON.stringify(value);

        default:
          return "";
      }
    }
  }, [data, nodeData]);

  const transform = `translate(${labelX}px,${labelY}px) translate(-50%, -50%)`;

  return (
    <>
      <BaseEdge id={id} path={edgePath} markerEnd={markerEnd} style={style} />
      {data.key && (
        <EdgeLabelRenderer>
          <div
            className="absolute rounded border bg-background px-1 text-foreground"
            style={{ transform }}
          >
            <pre className="text-xs">{label}</pre>
          </div>
        </EdgeLabelRenderer>
      )}
    </>
  );
}

/**
 * Chooses which of React Flow's edge path algorithms to use based on the provided
 * `type`.
 */
function getPath({
  type,
  sourceX,
  sourceY,
  targetX,
  targetY,
  sourcePosition,
  targetPosition,
}: {
  type: "bezier" | "smoothstep" | "step" | "straight";
  sourceX: number;
  sourceY: number;
  targetX: number;
  targetY: number;
  sourcePosition: Position;
  targetPosition: Position;
}) {
  switch (type) {
    case "bezier":
      return getBezierPath({
        sourceX,
        sourceY,
        targetX,
        targetY,
        sourcePosition,
        targetPosition,
      });

    case "smoothstep":
      return getSmoothStepPath({
        sourceX,
        sourceY,
        targetX,
        targetY,
        sourcePosition,
        targetPosition,
      });

    case "step":
      return getSmoothStepPath({
        sourceX,
        sourceY,
        targetX,
        targetY,
        sourcePosition,
        targetPosition,
        borderRadius: 0,
      });

    case "straight":
      return getStraightPath({
        sourceX,
        sourceY,
        targetX,
        targetY,
      });
  }
}


Data Edge

Installation

Additional type safety