check if edge is valid when connecting; add arrow marker to edges

koonpeng · koonpeng · commit b5ef3a72cb28 · 2025-07-31T04:15:59.000Z
Signed-off-by: Teo Koon Peng &lt;teokoonpeng@gmail.com&gt;
diff --git a/diagram-editor/frontend/add-operation.tsx b/diagram-editor/frontend/add-operation.tsx
@@ -11,6 +11,11 @@ import type { DiagramEditorNode } from './nodes';
 import {
   BufferAccessIcon,
   BufferIcon,
+  createOperationNode,
+  createScopeNode,
+  createSectionBufferNode,
+  createSectionInputNode,
+  createSectionOutputNode,
   ForkCloneIcon,
   ForkResultIcon,
   isOperationNode,
@@ -30,13 +35,6 @@ import {
   UnzipIcon,
 } from './nodes';
 import type { DiagramOperation } from './types/api';
-import {
-  createOperationNode,
-  createScopeNode,
-  createSectionBufferNode,
-  createSectionInputNode,
-  createSectionOutputNode,
-} from './utils/create-node';
 import { joinNamespaces, ROOT_NAMESPACE } from './utils/namespace';
 
 const StyledOperationButton = styled(Button)({
diff --git a/diagram-editor/frontend/diagram-editor.tsx b/diagram-editor/frontend/diagram-editor.tsx
@@ -39,6 +39,7 @@ import {
 import ExportDiagramDialog from './export-diagram-dialog';
 import { defaultEdgeData, EditEdgeForm, EditNodeForm } from './forms';
 import EditScopeForm from './forms/edit-scope-form';
+import { NodeManager } from './node-manager';
 import {
   type DiagramEditorNode,
   isOperationNode,
@@ -48,12 +49,11 @@ import {
 } from './nodes';
 import { useTemplates } from './templates-provider';
 import { autoLayout } from './utils/auto-layout';
-import { allowEdges as getAllowEdges } from './utils/connection';
+import { getValidEdgeTypes } from './utils/connection';
 import { exhaustiveCheck } from './utils/exhaustive-check';
+import { exportTemplate } from './utils/export-diagram';
 import { calculateScopeBounds, LAYOUT_OPTIONS } from './utils/layout';
 import { loadDiagramJson, loadEmpty, loadTemplate } from './utils/load-diagram';
-import { exportTemplate } from './utils/export-diagram';
-import { NodeManager } from './node-manager';
 
 const NonCapturingPopoverContainer = ({
   children,
@@ -527,14 +527,14 @@ function DiagramEditor() {
           closeAllPopovers();
         }}
         onConnect={(conn) => {
-          const sourceNode = nodes.find((n) => n.id === conn.source);
-          const targetNode = nodes.find((n) => n.id === conn.target);
+          const sourceNode = reactFlowInstance.current?.getNode(conn.source);
+          const targetNode = reactFlowInstance.current?.getNode(conn.target);
           if (!sourceNode || !targetNode) {
             throw new Error('cannot find source or target node');
           }
 
-          const allowedEdges = getAllowEdges(sourceNode, targetNode);
-          if (allowedEdges.length === 0) {
+          const validEdges = getValidEdgeTypes(sourceNode, targetNode);
+          if (validEdges.length === 0) {
             showErrorToast(
               `cannot connect "${sourceNode.type}" to "${targetNode.type}"`,
             );
@@ -545,13 +545,23 @@ function DiagramEditor() {
             addEdge(
               {
                 ...conn,
-                type: allowedEdges[0],
-                data: defaultEdgeData(allowedEdges[0]),
+                type: validEdges[0],
+                data: defaultEdgeData(validEdges[0]),
               },
               prev,
             ),
           );
         }}
+        isValidConnection={(conn) => {
+          const sourceNode = reactFlowInstance.current?.getNode(conn.source);
+          const targetNode = reactFlowInstance.current?.getNode(conn.target);
+          if (!sourceNode || !targetNode) {
+            throw new Error('cannot find source or target node');
+          }
+
+          const allowedEdges = getValidEdgeTypes(sourceNode, targetNode);
+          return allowedEdges.length > 0;
+        }}
         onReconnect={(oldEdge, newConnection) =>
           setEdges((prev) => reconnectEdge(oldEdge, newConnection, prev))
         }
@@ -678,7 +688,7 @@ function DiagramEditor() {
           {editingEdge && (
             <EditEdgeForm
               edge={editingEdge.edge}
-              allowedEdgeTypes={getAllowEdges(
+              allowedEdgeTypes={getValidEdgeTypes(
                 editingEdge.sourceNode,
                 editingEdge.targetNode,
               )}
diff --git a/diagram-editor/frontend/edges/create-edge.ts b/diagram-editor/frontend/edges/create-edge.ts
@@ -0,0 +1,135 @@
+import { MarkerType } from '@xyflow/react';
+import { v4 as uuidv4 } from 'uuid';
+import type { DiagramEditorEdge, EdgeData } from '.';
+
+function createBaseEdge(
+  source: string,
+  target: string,
+): Pick<DiagramEditorEdge, 'id' | 'source' | 'target' | 'markerEnd'> {
+  return {
+    id: uuidv4(),
+    source,
+    target,
+    markerEnd: {
+      type: MarkerType.ArrowClosed,
+      width: 24,
+      height: 24,
+    },
+  };
+}
+
+export function createDefaultEdge(
+  source: string,
+  target: string,
+): DiagramEditorEdge<'default'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'default',
+    data: {},
+  };
+}
+
+export function createUnzipEdge(
+  source: string,
+  target: string,
+  data: EdgeData<'unzip'>,
+): DiagramEditorEdge<'unzip'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'unzip',
+    data,
+  };
+}
+
+export function createForkResultOkEdge(
+  source: string,
+  target: string,
+): DiagramEditorEdge<'forkResultOk'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'forkResultOk',
+    data: {},
+  };
+}
+
+export function createForkResultErrEdge(
+  source: string,
+  target: string,
+): DiagramEditorEdge<'forkResultErr'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'forkResultErr',
+    data: {},
+  };
+}
+
+export function createSplitKeyEdge(
+  source: string,
+  target: string,
+  data: EdgeData<'splitKey'>,
+): DiagramEditorEdge<'splitKey'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'splitKey',
+    data,
+  };
+}
+
+export function createSplitSeqEdge(
+  source: string,
+  target: string,
+  data: EdgeData<'splitSeq'>,
+): DiagramEditorEdge<'splitSeq'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'splitSeq',
+    data,
+  };
+}
+
+export function createSplitRemainingEdge(
+  source: string,
+  target: string,
+): DiagramEditorEdge<'splitRemaining'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'splitRemaining',
+    data: {},
+  };
+}
+
+export function createBufferKeyEdge(
+  source: string,
+  target: string,
+  data: EdgeData<'bufferKey'>,
+): DiagramEditorEdge<'bufferKey'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'bufferKey',
+    data,
+  };
+}
+
+export function createBufferSeqEdge(
+  source: string,
+  target: string,
+  data: EdgeData<'bufferSeq'>,
+): DiagramEditorEdge<'bufferSeq'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'bufferSeq',
+    data,
+  };
+}
+
+export function createStreamOutEdge(
+  source: string,
+  target: string,
+  data: EdgeData<'streamOut'>,
+): DiagramEditorEdge<'streamOut'> {
+  return {
+    ...createBaseEdge(source, target),
+    type: 'streamOut',
+    data,
+  };
+}
diff --git a/diagram-editor/frontend/edges/index.ts b/diagram-editor/frontend/edges/index.ts
@@ -45,6 +45,30 @@ export type { SplitSeqEdge } from './split-seq-edge';
 export type { StreamOutEdge } from './stream-out-edge';
 export type { UnzipEdge } from './unzip-edge';
 
+export * from './create-edge';
+export type DefaultEdgeData = Record<string, never>;
+export type DefaultEdge = Edge<DefaultEdgeData, 'default'>;
+
+type EdgeMapping = {
+  default: { comp: DefaultEdge; data: DefaultEdgeData };
+  unzip: { comp: UnzipEdge; data: UnzipEdgeData };
+  forkResultOk: { comp: ForkResultOkEdge; data: ForkResultOkEdgeData };
+  forkResultErr: { comp: ForkResultErrEdge; data: ForkResultErrEdgeData };
+  splitKey: { comp: SplitKeyEdge; data: SplitKeyEdgeData };
+  splitSeq: { comp: SplitSeqEdge; data: SplitSeqEdgeData };
+  splitRemaining: { comp: SplitRemainingEdge; data: SplitRemainingEdgeData };
+  bufferKey: { comp: BufferKeyEdge; data: BufferKeyEdgeData };
+  bufferSeq: { comp: BufferSeqEdge; data: BufferSeqEdgeData };
+  streamOut: { comp: StreamOutEdge; data: StreamOutEdgeData };
+};
+
+export type EdgeTypes = keyof EdgeMapping;
+
+export type DiagramEditorEdge<T extends EdgeTypes = EdgeTypes> =
+  EdgeMapping[T]['comp'];
+
+export type EdgeData<T extends EdgeTypes = EdgeTypes> = EdgeMapping[T]['data'];
+
 export const EDGE_TYPES = {
   default: StepEdge,
   unzip: UnzipEdgeComp,
@@ -56,34 +80,4 @@ export const EDGE_TYPES = {
   bufferKey: BufferKeyEdgeComp,
   bufferSeq: BufferSeqEdgeComp,
   streamOut: StreamOutEdgeComp,
-};
-
-export type EdgeTypes = keyof typeof EDGE_TYPES;
-
-export type DefaultEdgeData = Record<string, never>;
-export type DefaultEdge = Edge<DefaultEdgeData, 'default'>;
-
-export type EdgeData = {
-  default: DefaultEdgeData;
-  unzip: UnzipEdgeData;
-  forkResultOk: ForkResultOkEdgeData;
-  forkResultErr: ForkResultErrEdgeData;
-  splitKey: SplitKeyEdgeData;
-  splitSeq: SplitSeqEdgeData;
-  splitRemaining: SplitRemainingEdgeData;
-  bufferKey: BufferKeyEdgeData;
-  bufferSeq: BufferSeqEdgeData;
-  streamOut: StreamOutEdgeData;
-};
-
-export type DiagramEditorEdge =
-  | DefaultEdge
-  | UnzipEdge
-  | ForkResultOkEdge
-  | ForkResultErrEdge
-  | SplitKeyEdge
-  | SplitSeqEdge
-  | SplitRemainingEdge
-  | BufferKeyEdge
-  | BufferSeqEdge
-  | StreamOutEdge;
+} satisfies Record<EdgeTypes, unknown>;
diff --git a/diagram-editor/frontend/nodes/create-node.ts b/diagram-editor/frontend/nodes/create-node.ts
@@ -1,5 +1,8 @@
 import type { XYPosition } from '@xyflow/react';
 import { v4 as uuidv4 } from 'uuid';
+import type { DiagramOperation, NextOperation } from '../types/api';
+import { calculateScopeBounds, LAYOUT_OPTIONS } from '../utils/layout';
+import { joinNamespaces } from '../utils/namespace';
 import {
   type DiagramEditorNode,
   type OperationNode,
@@ -8,10 +11,7 @@ import {
   type SectionOutputNode,
   START_ID,
   TERMINATE_ID,
-} from '../nodes';
-import type { DiagramOperation, NextOperation } from '../types/api';
-import { calculateScopeBounds, LAYOUT_OPTIONS } from './layout';
-import { joinNamespaces } from './namespace';
+} from '.';
 
 export function createStartNode(
   namespace: string,
diff --git a/diagram-editor/frontend/nodes/index.ts b/diagram-editor/frontend/nodes/index.ts
@@ -22,6 +22,7 @@ import TerminateNodeComp from './terminate-node';
 import TransformNodeComp from './transform-node';
 import UnzipNodeComp from './unzip-node';
 
+export * from './create-node';
 export * from './icons';
 export type {
   SectionBufferData,
diff --git a/diagram-editor/frontend/utils/connection.ts b/diagram-editor/frontend/utils/connection.ts
@@ -55,15 +55,15 @@ function setIntersection<T>(a: Set<T>, b: Set<T>): Set<T> {
   return intersection;
 }
 
-export function allowEdges(
-  source: DiagramEditorNode,
-  target: DiagramEditorNode,
+export function getValidEdgeTypes(
+  sourceNode: DiagramEditorNode,
+  targetNode: DiagramEditorNode,
 ): EdgeTypes[] {
-  const allowed_output_edges = source.type
-    ? ALLOWED_OUTPUT_EDGES[source.type]
+  const allowedOutputEdges = sourceNode.type
+    ? ALLOWED_OUTPUT_EDGES[sourceNode.type]
     : new Set([]);
-  const allowed_input_edges = target.type
-    ? ALLOWED_INPUT_EDGES[target.type]
+  const allowedInputEdges = targetNode.type
+    ? ALLOWED_INPUT_EDGES[targetNode.type]
     : new Set([]);
-  return Array.from(setIntersection(allowed_output_edges, allowed_input_edges));
+  return Array.from(setIntersection(allowedOutputEdges, allowedInputEdges));
 }
diff --git a/diagram-editor/frontend/utils/export-diagram.test.ts b/diagram-editor/frontend/utils/export-diagram.test.ts
@@ -1,13 +1,14 @@
 import { v4 as uuidv4 } from 'uuid';
 import type { DiagramEditorEdge } from '../edges';
 import { NodeManager } from '../node-manager';
-import { type DiagramEditorNode, START_ID } from '../nodes';
 import {
   createOperationNode,
   createSectionBufferNode,
   createSectionInputNode,
   createSectionOutputNode,
-} from './create-node';
+  type DiagramEditorNode,
+  START_ID,
+} from '../nodes';
 import { exportDiagram, exportTemplate } from './export-diagram';
 import { loadDiagramJson } from './load-diagram';
 import { joinNamespaces, ROOT_NAMESPACE } from './namespace';
diff --git a/diagram-editor/frontend/utils/load-diagram.ts b/diagram-editor/frontend/utils/load-diagram.ts
diff --git a/diagram-editor/frontend/utils/operation.ts b/diagram-editor/frontend/utils/operation.ts
