feat(sdk): edge template prop (#28)

* feat(sdk): expose edgeTemplates prop for custom edge renderers (WB-220)

* feat(demo): custom edge example via edgeTemplates (WB-220)

* fix(demo): reflect selection/hover state in custom DashedEdge (WB-220)

The demo custom edge painted a fixed stroke, so it didn't highlight on
select/hover like the built-in edges — a misleading reference. Delegate
styling to the SDK's exported useLabelEdgeHover (the same hook LabelEdge
uses), keeping only strokeDasharray on top. Selected/hover now match the
built-in edge while the edge stays visually dashed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs(sdk): document customizing edge selection (WB-220)

* docs(sdk): address review findings on edgeTemplates (WB-220)

* docs(sdk): clarify edge-types naming in DiagramContainer (WB-220)

Rename the vague `resolvedEdgeTypes` local to `baseEdgeTypes` and correct
the comment: the value is the built-in `labelEdge` renderer plus app-wide
`edgeTemplates` from <WorkflowBuilder.Root>, which the per-mount `edgeTypes`
prop overrides — not plugin- or ELK-injected edges.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* docs(sdk): refine edgeTemplates docs, trim changeset

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Piotr Blaszczyk <piotr.blaszczyk@synergycodes.com>

Author: 3 people

.changeset/wb-220-edge-templates.md

@@ -0,0 +1,5 @@
+---
+'@workflowbuilder/sdk': minor
+---
+
+feat: add `edgeTemplates` prop on `<WorkflowBuilder.Root>` for custom edge renderers. Pass a `{ [edgeType]: Component }` map of components taking ReactFlow's `EdgeProps`; edges whose `type` matches a key render with your component, and unregistered types fall back to the built-in `labelEdge`. Also exports the `WorkflowBuilderEdgeTemplates` type.

apps/demo/src/app/app.tsx

@@ -1,7 +1,9 @@
 import { WorkflowBuilder } from '@workflowbuilder/sdk';
+import type { WorkflowBuilderEdgeTemplates, WorkflowBuilderNodeTemplates } from '@workflowbuilder/sdk';
 
 import '@workflowbuilder/sdk/style.css';
 
+import { DashedEdge } from './components/dashed-edge/dashed-edge';
 import { MultiPortNodeTemplate } from './components/multi-port-node/multi-port-node-template';
 import { demoPaletteItems } from './data/palette';
 import { demoTemplates } from './data/templates';
@@ -18,14 +20,23 @@ import { plugin as undoRedoPlugin } from './plugins/undo-redo/plugin-exports';
 import { plugin as validationPlugin } from './plugins/validation/plugin-exports';
 import { plugin as widgetsPlugin } from './plugins/widgets/plugin-exports';
 
+// Declared at module scope so the references stay stable across renders, as
+// `<WorkflowBuilder.Root>` requires for nodeTemplates / edgeTemplates.
+const nodeTemplates = {
+  'multi-port': MultiPortNodeTemplate,
+} satisfies WorkflowBuilderNodeTemplates;
+
+const edgeTemplates = {
+  dashed: DashedEdge,
+} satisfies WorkflowBuilderEdgeTemplates;
+
 export function App() {
   return (
     <WorkflowBuilder.Root
       name="demo"
       nodeTypes={demoPaletteItems}
-      nodeTemplates={{
-        'multi-port': MultiPortNodeTemplate,
-      }}
+      nodeTemplates={nodeTemplates}
+      edgeTemplates={edgeTemplates}
       diagramTemplates={demoTemplates}
       plugins={[
         demoPlugin,

apps/demo/src/app/components/dashed-edge/dashed-edge.tsx

@@ -0,0 +1,41 @@
+import { EnhancedBaseEdge, useLabelEdgeHover } from '@workflowbuilder/sdk';
+import { type EdgeProps, getSmoothStepPath } from '@xyflow/react';
+
+/**
+ * Demo custom edge. Registered on `<WorkflowBuilder.Root edgeTemplates>` under
+ * the `'dashed'` key, so any edge with `type: 'dashed'` renders dashed instead
+ * of the built-in `labelEdge`.
+ *
+ * Mirrors the `multi-port` node-template example: a consumer component that
+ * takes ReactFlow's `EdgeProps` directly (no SDK adapter) and reuses the
+ * exported `EnhancedBaseEdge` for a wide hit target. Selection and hover are
+ * delegated to the SDK's `useLabelEdgeHover` so this edge highlights exactly
+ * like the built-in one — it just keeps a dashed stroke on top.
+ *
+ * This is a minimal styling example, not a full `labelEdge` replacement. It
+ * deliberately skips two things the built-in does: the self-connecting loop
+ * when `source === target` (a self-loop drawn with this type collapses to a
+ * degenerate path), and rendering `data.label` / `data.icon`. If you need
+ * either, branch on `source === target` and delegate to the exported
+ * `SelfConnectingEdge`, and render your own label. See the SDK README,
+ * "Custom edges and selection".
+ *
+ * To diverge from the built-in selection look, add a `selected` branch to the
+ * `style` below, or restyle every edge globally via the `--ax-public-edge-color-select`
+ * CSS variable. See the SDK README, "Custom edges and selection".
+ */
+export function DashedEdge({
+  id,
+  sourceX,
+  sourceY,
+  targetX,
+  targetY,
+  sourcePosition,
+  targetPosition,
+  selected,
+}: EdgeProps) {
+  const { style } = useLabelEdgeHover({ id, isSelected: selected });
+  const [path] = getSmoothStepPath({ sourceX, sourceY, targetX, targetY, sourcePosition, targetPosition });
+
+  return <EnhancedBaseEdge id={id} path={path} style={{ ...style, strokeDasharray: '6 4' }} />;
+}

apps/demo/src/app/data/templates/simple-flow.ts

@@ -299,7 +299,9 @@ const defaultDiagram: DiagramModel = {
         target: 'da47caa9-c695-47bb-be52-b30bb8a6be6d',
         targetHandle: 'target',
         zIndex: 0,
-        type: 'labelEdge',
+        // Custom edge type registered via `edgeTemplates` in app.tsx — renders
+        // the demo's DashedEdge instead of the built-in labelEdge.
+        type: 'dashed',
         id: 'xy-edge__440ccd46-0f50-4e35-ae74-64fee988a4f6440ccd46-0f50-4e35-ae74-64fee988a4f6:source-da47caa9-c695-47bb-be52-b30bb8a6be6dda47caa9-c695-47bb-be52-b30bb8a6be6d:target',
       },
       {

packages/sdk/README.md

@@ -89,18 +89,19 @@ If you omit `<WorkflowBuilder.TopBar />`, use [`useWorkflowBuilderActions()`](ht
 
 ## `<WorkflowBuilder.Root>` props
 
-| Prop               | Type                            | Description                                                                                                                                                                    |
-| ------------------ | ------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| `nodeTypes`        | `PaletteItemOrGroup[]`          | Node type definitions. Appear in the palette and drive validation. **Must be a stable reference** — declare at module scope or memoize.                                        |
-| `nodeTemplates`    | `WorkflowBuilderNodeTemplates`  | Per-node-type custom renderers. Map of `data.type` → React component, overriding the default node renderer for that type. **Stable reference required** (same as `nodeTypes`). |
-| `diagramTemplates` | `TemplateModel[]`               | Diagram templates available in the template selector. **Stable reference required** (same as `nodeTypes`).                                                                     |
-| `plugins`          | `WorkflowBuilderPlugin[]`       | Functions registering decorators. Synchronous, executed once.                                                                                                                  |
-| `jsonForm`         | `WorkflowBuilderJsonFormConfig` | Custom JsonForms renderers, cells, translations.                                                                                                                               |
-| `integration`      | `WorkflowBuilderIntegration`    | Data source / sink. Defaults to `localStorage`.                                                                                                                                |
-| `name`             | `string`                        | Workflow name shown in the header.                                                                                                                                             |
-| `layoutDirection`  | `'DOWN' \| 'RIGHT'`             | Initial flow direction.                                                                                                                                                        |
-| `initialNodes`     | `WorkflowBuilderNode[]`         | Starting diagram nodes.                                                                                                                                                        |
-| `initialEdges`     | `WorkflowBuilderEdge[]`         | Starting diagram edges.                                                                                                                                                        |
+| Prop               | Type                            | Description                                                                                                                                                                                                                                    |
+| ------------------ | ------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `nodeTypes`        | `PaletteItemOrGroup[]`          | Node type definitions. Appear in the palette and drive validation. **Must be a stable reference** — declare at module scope or memoize.                                                                                                        |
+| `nodeTemplates`    | `WorkflowBuilderNodeTemplates`  | Per-node-type custom renderers. Map of `data.type` → React component, overriding the default node renderer for that type. **Stable reference required** (same as `nodeTypes`).                                                                 |
+| `edgeTemplates`    | `WorkflowBuilderEdgeTemplates`  | Per-edge-type custom renderers. Map of `edge.type` → React component taking ReactFlow `EdgeProps`, overriding the built-in `labelEdge`. Unregistered types fall back to the default edge. **Stable reference required** (same as `nodeTypes`). |
+| `diagramTemplates` | `TemplateModel[]`               | Diagram templates available in the template selector. **Stable reference required** (same as `nodeTypes`).                                                                                                                                     |
+| `plugins`          | `WorkflowBuilderPlugin[]`       | Functions registering decorators. Synchronous, executed once.                                                                                                                                                                                  |
+| `jsonForm`         | `WorkflowBuilderJsonFormConfig` | Custom JsonForms renderers, cells, translations.                                                                                                                                                                                               |
+| `integration`      | `WorkflowBuilderIntegration`    | Data source / sink. Defaults to `localStorage`.                                                                                                                                                                                                |
+| `name`             | `string`                        | Workflow name shown in the header.                                                                                                                                                                                                             |
+| `layoutDirection`  | `'DOWN' \| 'RIGHT'`             | Initial flow direction.                                                                                                                                                                                                                        |
+| `initialNodes`     | `WorkflowBuilderNode[]`         | Starting diagram nodes.                                                                                                                                                                                                                        |
+| `initialEdges`     | `WorkflowBuilderEdge[]`         | Starting diagram edges.                                                                                                                                                                                                                        |
 
 Full reference (every public type, hook, and helper): <https://www.workflowbuilder.io/docs/api/core/workflowbuilder/>.
 

packages/sdk/src/data/edge-templates.ts

@@ -0,0 +1,18 @@
+import type { EdgeProps } from '@xyflow/react';
+import type { ComponentType } from 'react';
+
+import type { WorkflowBuilderEdge } from '../node/node-data';
+
+type EdgeTemplateComponent = ComponentType<EdgeProps<WorkflowBuilderEdge>>;
+
+const EMPTY: Readonly<Record<string, EdgeTemplateComponent>> = Object.freeze({});
+
+let customEdgeTemplates: Record<string, EdgeTemplateComponent> = EMPTY;
+
+export function setCustomEdgeTemplates(templates: Record<string, EdgeTemplateComponent> | null): void {
+  customEdgeTemplates = templates ?? EMPTY;
+}
+
+export function getCustomEdgeTemplates(): Record<string, EdgeTemplateComponent> {
+  return customEdgeTemplates;
+}

packages/sdk/src/features/diagram/diagram.tsx

@@ -27,8 +27,8 @@ import { useDeleteConfirmation } from '../modals/delete-confirmation/use-delete-
 import { withOptionalComponentPlugins } from '../plugins-core/adapters/adapter-components';
 import { deleteKeyCode } from './const';
 import { SNAP_GRID, SNAP_IS_ACTIVE } from './diagram.const';
-import { LabelEdge } from './edges/label-edge/label-edge';
 import { TemporaryEdge } from './edges/temporary-edge/temporary-edge';
+import { useEdgeTypes } from './hooks/use-edge-types';
 import { useNodeTypes } from './hooks/use-node-types';
 import { callNodeChangedListeners } from './listeners/node-changed-listeners';
 import { callNodeDragStartListeners } from './listeners/node-drag-start-listeners';
@@ -42,7 +42,12 @@ import { diagramStateSelector } from './selectors';
  * @category Components
  */
 export type DiagramContainerProps = {
-  /** Extra edge types forwarded to ReactFlow alongside the built-in `'labelEdge'`. */
+  /**
+   * Extra edge types forwarded to ReactFlow alongside the built-in `'labelEdge'`
+   * and any Root-level `edgeTemplates`. Merged last, so a key here intentionally
+   * overrides those (this is the direct-mount escape hatch, hence no collision
+   * warning); prefer `<WorkflowBuilder.Root edgeTemplates>` for app-wide edges.
+   */
   edgeTypes?: EdgeTypes;
 };
 
@@ -148,7 +153,11 @@ function DiagramContainerComponent({ edgeTypes = {} }: DiagramContainerProps) {
     [onSelectionChange],
   );
 
-  const diagramEdgeTypes = useMemo(() => ({ labelEdge: LabelEdge, ...edgeTypes }), [edgeTypes]);
+  // Built-in edge renderers (`labelEdge`) plus any app-wide `edgeTemplates`
+  // passed to `<WorkflowBuilder.Root>`. The local `edgeTypes` prop (direct
+  // DiagramContainer mount) is merged last so a per-mount override still wins.
+  const baseEdgeTypes = useEdgeTypes();
+  const diagramEdgeTypes = useMemo(() => ({ ...baseEdgeTypes, ...edgeTypes }), [baseEdgeTypes, edgeTypes]);
 
   const onBeforeDelete: OnBeforeDelete<WorkflowBuilderNode, WorkflowBuilderEdge> = useCallback(
     async ({ nodes, edges }) => {

packages/sdk/src/features/diagram/hooks/use-edge-types.spec.ts

@@ -0,0 +1,65 @@
+import { renderHook } from '@testing-library/react';
+import { afterEach, describe, expect, it, vi } from 'vitest';
+
+import { setCustomEdgeTemplates } from '../../../data/edge-templates';
+
+// Mock the built-in default edge so the test doesn't pull the full edge
+// rendering stack (xyflow path math, CSS side effects).
+vi.mock('../edges/label-edge/label-edge', () => ({ LabelEdge: () => null }));
+
+const { useEdgeTypes } = await import('./use-edge-types');
+
+function Noop() {
+  return null;
+}
+
+describe('useEdgeTypes', () => {
+  afterEach(() => {
+    setCustomEdgeTemplates(null);
+    vi.restoreAllMocks();
+  });
+
+  it('registers the built-in labelEdge by default', () => {
+    const { result } = renderHook(() => useEdgeTypes());
+
+    expect(result.current['labelEdge']).toBeDefined();
+  });
+
+  it('registers a custom edge template under its own key, unwrapped', () => {
+    setCustomEdgeTemplates({ animated: Noop });
+
+    const { result } = renderHook(() => useEdgeTypes());
+
+    // No adapter: the consumer's component is registered as-is (identity),
+    // unlike node templates which are wrapped to inject computed props.
+    expect(result.current['animated']).toBe(Noop);
+    expect(result.current['labelEdge']).toBeDefined();
+  });
+
+  it('warns in dev when a custom key collides with the built-in labelEdge', () => {
+    const spy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    setCustomEdgeTemplates({ labelEdge: Noop });
+
+    renderHook(() => useEdgeTypes());
+
+    expect(spy).toHaveBeenCalledWith(expect.stringContaining('"labelEdge"'));
+  });
+
+  it('does not warn for non-colliding custom edge keys', () => {
+    const spy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    setCustomEdgeTemplates({ animated: Noop });
+
+    renderHook(() => useEdgeTypes());
+
+    expect(spy).not.toHaveBeenCalled();
+  });
+
+  it('returns a stable reference across re-renders when no custom templates exist', () => {
+    const { result, rerender } = renderHook(() => useEdgeTypes());
+    const first = result.current;
+
+    rerender();
+
+    expect(result.current).toBe(first);
+  });
+});

packages/sdk/src/features/diagram/hooks/use-edge-types.tsx

@@ -0,0 +1,43 @@
+import type { EdgeTypes } from '@xyflow/react';
+import { useMemo } from 'react';
+
+import { getCustomEdgeTemplates } from '../../../data/edge-templates';
+import { LabelEdge } from '../edges/label-edge/label-edge';
+
+const BUILT_IN_KEYS: ReadonlySet<string> = new Set<string>(['labelEdge']);
+
+/**
+ * Resolves the ReactFlow `edgeTypes` map: the built-in `'labelEdge'` default
+ * merged with any consumer-provided `edgeTemplates`. Mirrors {@link useNodeTypes}
+ * but without an adapter — edge templates take ReactFlow's `EdgeProps` directly
+ * (the built-in edges do too), so there are no computed props to inject and the
+ * consumer's component drops straight into the map.
+ *
+ * An edge whose `type` matches no key falls back to ReactFlow's default edge.
+ */
+export function useEdgeTypes(): EdgeTypes {
+  // Read outside useMemo so the dep stays referentially stable across renders.
+  // When no consumer templates exist, getCustomEdgeTemplates() returns the same
+  // frozen EMPTY object every call (see ../../../data/edge-templates), which is
+  // what keeps this memo from handing ReactFlow a fresh edgeTypes object — and
+  // emitting its "it looks like you've created a new edgeTypes object" warning —
+  // on every render.
+  const custom = getCustomEdgeTemplates();
+
+  return useMemo<EdgeTypes>(() => {
+    if (import.meta.env.DEV) {
+      for (const key of Object.keys(custom)) {
+        if (BUILT_IN_KEYS.has(key)) {
+          console.warn(
+            `[workflow-builder] edgeTemplates key "${key}" overrides a built-in renderer. Pick a unique key unless the override is intentional.`,
+          );
+        }
+      }
+    }
+
+    return {
+      labelEdge: LabelEdge,
+      ...custom,
+    };
+  }, [custom]);
+}

packages/sdk/src/index.ts

@@ -55,6 +55,7 @@ export type {
   WorkflowBuilderIntegration,
   WorkflowBuilderJsonFormConfig,
   WorkflowBuilderNodeTemplates,
+  WorkflowBuilderEdgeTemplates,
 } from './workflow-builder-root';
 
 // =============================================================================