feat: Add custom catch container node

packages/serverless-workflow-diagram-editor/src/react-flow/diagram/diagramBuilder.ts

@@ -16,7 +16,7 @@
 
 import * as RF from "@xyflow/react";
 import { buildFlatGraph } from "../../core";
-import { BaseNodeData, ReactFlowNodeTypes } from "../nodes/Nodes";
+import { BaseNodeData, CATCH_CONTAINER_NODE_TYPE, ReactFlowNodeTypes } from "../nodes/Nodes";
 import { BaseEdgeData, EdgeTypes } from "../edges/Edges";
 import * as sdk from "@serverlessworkflow/sdk";
 import { applyAutoLayout, DEFAULT_NODE_SIZE } from "./autoLayout";
@@ -46,15 +46,45 @@ export function edgeSourceAndTargetExist(
   return nodeIdSet.has(graphEdge.sourceId) && nodeIdSet.has(graphEdge.targetId);
 }
 
-function buildReactFlowNode(graphNode: sdk.FlatGraph | sdk.FlatGraphNode): RF.Node<BaseNodeData> {
+/* Return ids of catch nodes that are containers (have child nodes with parentIds pointing to them) */
+export function getCatchContainerNodeIds(graph: sdk.FlatGraph): Set<string> {
+  const parentIds = new Set<string>();
+  for (const node of graph.nodes) {
+    if (node.parentId) {
+      parentIds.add(node.parentId);
+    }
+  }
+
+  const containerIds = new Set<string>();
+  for (const node of graph.nodes) {
+    if (node.type === sdk.GraphNodeType.Catch && parentIds.has(node.id)) {
+      containerIds.add(node.id);
+    }
+  }
+
+  return containerIds;
+}
+
+function resolveNodeType(graphNode: sdk.FlatGraphNode, catchContainerIds: Set<string>): string {
+  if (graphNode.type === sdk.GraphNodeType.Catch && catchContainerIds.has(graphNode.id)) {
+    return CATCH_CONTAINER_NODE_TYPE;
+  }
+  return graphNode.type;
+}
+
+function buildReactFlowNode(
+  graphNode: sdk.FlatGraphNode,
+  catchContainerIds: Set<string>,
+): RF.Node<BaseNodeData> {
+  const type = resolveNodeType(graphNode, catchContainerIds);
   // There is no corresponding react flow component implemented
-  if (!Object.keys(ReactFlowNodeTypes).includes(graphNode.type)) {
-    throw new Error(`Unsupported GraphNodeType: ${graphNode.type}!`);
+  if (!Object.keys(ReactFlowNodeTypes).includes(type)) {
+    throw new Error(`Unsupported React flow node type: ${type}!`);
   }
 
   return {
     id: graphNode.id,
-    type: graphNode.type,
+    type,
     data: {
       label: graphNode.label ?? "",
       ...(graphNode.task !== undefined && { task: structuredClone(graphNode.task) }),
@@ -89,11 +119,12 @@ export function buildDiagramElements(model: sdk.Specification.Workflow | null):
 
   if (model) {
     const graph = buildFlatGraph(model);
+    const catchContainerIds = getCatchContainerNodeIds(graph);
 
     graph.nodes.forEach((graphNode) => {
       // TODO: only nodes on root level are supported for now
       if (graphNode.parentId === "root") {
-        nodes.push(buildReactFlowNode(graphNode));
+        nodes.push(buildReactFlowNode(graphNode, catchContainerIds));
       }
     });
 

packages/serverless-workflow-diagram-editor/src/react-flow/nodes/Nodes.tsx

@@ -21,7 +21,9 @@ import { type LeafNodeType, taskNodeConfigMap } from "./taskNodeConfig";
 import { Info } from "lucide-react";
 import { getCallSubType, getListenSubType, getRunSubType } from "../../core";
 
-// Node types must match sdk GraphNodeType enum
+export const CATCH_CONTAINER_NODE_TYPE = "catch-container";
+/* Node types are primarily keyed by the sdk GraphNodeType enum, with custom
+   React Flow-only node types such as catch-container added where needed. */
 export const ReactFlowNodeTypes: RF.NodeTypes = {
   [GraphNodeType.Start]: StartNode,
   [GraphNodeType.End]: EndNode,
@@ -38,6 +40,7 @@ export const ReactFlowNodeTypes: RF.NodeTypes = {
   [GraphNodeType.TryCatch]: TryCatchNode,
   [GraphNodeType.Try]: TryNode,
   [GraphNodeType.Catch]: CatchNode,
+  [CATCH_CONTAINER_NODE_TYPE]: CatchContainerNode,
   [GraphNodeType.Wait]: WaitNode,
 };
 
@@ -250,6 +253,18 @@ export function CatchNode({ id, data, selected, type }: RF.NodeProps<CatchNodeTy
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
 }
 
+/* catch container node */
+export type CatchContainerNodeType = RF.Node<BaseNodeData, typeof CATCH_CONTAINER_NODE_TYPE>;
+export function CatchContainerNode({
+  id,
+  data,
+  selected,
+  type,
+}: RF.NodeProps<CatchContainerNodeType>) {
+  // TODO: This component is just a placeholder
+  return <PlaceholderContent id={id} data={data} selected={selected} type={type} />;
+}
+
 /* wait leaf node */
 export type WaitNodeType = RF.Node<BaseNodeData<Specification.WaitTask>, typeof GraphNodeType.Wait>;
 export function WaitNode({ id, data, selected, type }: RF.NodeProps<WaitNodeType>) {

packages/serverless-workflow-diagram-editor/tests/core/graph.test.ts

@@ -15,25 +15,9 @@
  */
 
 import { describe, it, expect } from "vitest";
-import { FlatGraph, FlatGraphNode, GraphNodeType } from "@serverlessworkflow/sdk";
+import { FlatGraphNode, GraphNodeType } from "@serverlessworkflow/sdk";
 import { getNodesByType, fixNodesConnections } from "../../src/core/graph";
-
-function createFlatGraph(
-  nodes: FlatGraphNode[],
-  edges: Array<{ id: string; sourceId: string; targetId: string; label: string }>,
-): FlatGraph {
-  const entryNode = nodes.find((n) => n.type === GraphNodeType.Entry);
-  const exitNode = nodes.find((n) => n.type === GraphNodeType.Exit);
-
-  return {
-    id: "root",
-    type: GraphNodeType.Do,
-    nodes,
-    edges,
-    entryNode,
-    exitNode,
-  } as FlatGraph;
-}
+import { createFlatGraph } from "../test-utils/graph-helpers";
 
 describe("graph utils", () => {
   describe("getNodesByType", () => {

packages/serverless-workflow-diagram-editor/tests/react-flow/diagram/diagramBuilder.test.ts

@@ -16,18 +16,20 @@
 
 import { describe, it, expect, beforeAll } from "vitest";
 import * as RF from "@xyflow/react";
-import { GraphNodeType } from "@serverlessworkflow/sdk";
+import { FlatGraphNode, GraphNodeType } from "@serverlessworkflow/sdk";
 import {
   getEdgeType,
   edgeSourceAndTargetExist,
   buildDiagramElements,
+  getCatchContainerNodeIds,
 } from "../../../src/react-flow/diagram/diagramBuilder";
 import { EdgeTypes } from "../../../src/react-flow/edges/Edges";
 import { parseWorkflow } from "../../../src/core";
 import {
   BASIC_VALID_WORKFLOW_JSON,
   BASIC_VALID_WORKFLOW_JSON_TASKS,
 } from "../../fixtures/workflows";
+import { createFlatGraph } from "../../test-utils/graph-helpers";
 
 // Type alias for diagram elements to reduce verbosity
 type DiagramElements = ReturnType<typeof buildDiagramElements>;
@@ -363,4 +365,139 @@ describe("diagramBuilder", () => {
       });
     });
   });
+
+  describe("getCatchContainerNodeIds", () => {
+    it("should not include leaf catch nodes", () => {
+      const sdkGraph = createFlatGraph(
+        [
+          {
+            id: "start",
+            type: GraphNodeType.Start,
+          } as FlatGraphNode,
+          {
+            id: "end",
+            type: GraphNodeType.End,
+          } as FlatGraphNode,
+          {
+            id: "/do/0/CatchError",
+            type: GraphNodeType.Catch,
+            label: "CatchError",
+          } as FlatGraphNode,
+        ],
+        [],
+      );
+
+      const result = getCatchContainerNodeIds(sdkGraph);
+      expect(result.size).toBe(0);
+    });
+
+    it("should include catch nodes that have children", () => {
+      const sdkGraph = createFlatGraph(
+        [
+          {
+            id: "start",
+            type: GraphNodeType.Start,
+          } as FlatGraphNode,
+          {
+            id: "end",
+            type: GraphNodeType.End,
+          } as FlatGraphNode,
+          {
+            id: "/do/0/CatchContainer",
+            type: GraphNodeType.Catch,
+            label: "CatchContainer",
+          } as FlatGraphNode,
+          {
+            id: "/do/0/CatchContainer-entry-node",
+            type: GraphNodeType.Entry,
+            parentId: "/do/0/CatchContainer",
+          } as FlatGraphNode,
+          {
+            id: "/do/0/CatchContainer-exit-node",
+            type: GraphNodeType.Exit,
+            parentId: "/do/0/CatchContainer",
+          } as FlatGraphNode,
+          {
+            id: "/do/0/CatchContainer/do/0/step1",
+            type: GraphNodeType.Set,
+            label: "step1",
+            parentId: "/do/0/CatchContainer",
+          } as FlatGraphNode,
+        ],
+        [],
+      );
+
+      const result = getCatchContainerNodeIds(sdkGraph);
+
+      expect(result.has("/do/0/CatchContainer")).toBe(true);
+      expect(result.size).toBe(1);
+    });
+
+    it("should not include non-catch nodes", () => {
+      const sdkGraph = createFlatGraph(
+        [
+          {
+            id: "start",
+            type: GraphNodeType.Start,
+          } as FlatGraphNode,
+          {
+            id: "end",
+            type: GraphNodeType.End,
+          } as FlatGraphNode,
+          {
+            id: "/do/0/step1",
+            type: GraphNodeType.Set,
+            label: "step1",
+          } as FlatGraphNode,
+          {
+            id: "/do/0/step2",
+            type: GraphNodeType.Call,
+            label: "step2",
+          } as FlatGraphNode,
+        ],
+        [],
+      );
+
+      const result = getCatchContainerNodeIds(sdkGraph);
+      expect(result.size).toBe(0);
+    });
+
+    it("should include nested catch containers inside a parent container", () => {
+      const sdkGraph = createFlatGraph(
+        [
+          {
+            id: "start",
+            type: GraphNodeType.Start,
+          } as FlatGraphNode,
+          {
+            id: "end",
+            type: GraphNodeType.End,
+          } as FlatGraphNode,
+          {
+            id: "/do/0/TryCatch",
+            type: GraphNodeType.TryCatch,
+            label: "TryCatch",
+          } as FlatGraphNode,
+          {
+            id: "/do/0/TryCatch/catch",
+            type: GraphNodeType.Catch,
+            label: "catch",
+            parentId: "/do/0/TryCatch",
+          } as FlatGraphNode,
+          {
+            id: "/do/0/TryCatch/catch/do/0/recover",
+            type: GraphNodeType.Set,
+            label: "recover",
+            parentId: "/do/0/TryCatch/catch",
+          } as FlatGraphNode,
+        ],
+        [],
+      );
+
+      const result = getCatchContainerNodeIds(sdkGraph);
+
+      expect(result.has("/do/0/TryCatch/catch")).toBe(true);
+      expect(result.has("/do/0/TryCatch")).toBe(false);
+    });
+  });
 });

packages/serverless-workflow-diagram-editor/tests/test-utils/graph-helpers.ts

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2021-Present The Serverless Workflow Specification Authors
+ *
+ * Licensed 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 { FlatGraph, FlatGraphNode, GraphNodeType } from "@serverlessworkflow/sdk";
+
+export function createFlatGraph(
+  nodes: FlatGraphNode[],
+  edges: Array<{ id: string; sourceId: string; targetId: string; label: string }>,
+): FlatGraph {
+  const entryNode = nodes.find((n) => n.type === GraphNodeType.Entry);
+  const exitNode = nodes.find((n) => n.type === GraphNodeType.Exit);
+
+  return {
+    id: "root",
+    type: GraphNodeType.Do,
+    nodes,
+    edges,
+    entryNode,
+    exitNode,
+  } as FlatGraph;
+}

packages/serverless-workflow-diagram-editor/tests/test-utils/index.ts

@@ -16,3 +16,4 @@
 
 export { renderWithProviders } from "./render-helpers";
 export { t } from "./translation-helpers";
+export { createFlatGraph } from "./graph-helpers";