diff --git a/packages/serverless-workflow-diagram-editor/src/core/index.ts b/packages/serverless-workflow-diagram-editor/src/core/index.ts
index 0bbbd0e..e4b9e03 100644
--- a/packages/serverless-workflow-diagram-editor/src/core/index.ts
+++ b/packages/serverless-workflow-diagram-editor/src/core/index.ts
@@ -16,3 +16,4 @@
 
 export * from "./workflowSdk";
 export * from "./graph";
+export * from "./taskSubType";
diff --git a/packages/serverless-workflow-diagram-editor/src/core/taskSubType.ts b/packages/serverless-workflow-diagram-editor/src/core/taskSubType.ts
new file mode 100644
index 0000000..396431f
--- /dev/null
+++ b/packages/serverless-workflow-diagram-editor/src/core/taskSubType.ts
@@ -0,0 +1,33 @@
+/*
+ * 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 type { Specification } from "@serverlessworkflow/sdk";
+
+function getFirstKey(obj: unknown): string | undefined {
+  return obj && typeof obj === "object" && !Array.isArray(obj) ? Object.keys(obj)[0] : undefined;
+}
+
+export function getRunSubType(task: Specification.RunTask): string | undefined {
+  return getFirstKey(task.run);
+}
+
+export function getListenSubType(task: Specification.ListenTask): string | undefined {
+  return getFirstKey(task.listen?.to);
+}
+
+export function getCallSubType(task: Specification.CallTask): string | undefined {
+  return typeof task.call === "string" ? task.call : undefined;
+}
diff --git a/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/Diagram.css b/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/Diagram.css
index b9deba5..04188fd 100644
--- a/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/Diagram.css
+++ b/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/Diagram.css
@@ -198,22 +198,25 @@
     @apply dec:flex 
       dec:items-center 
       dec:gap-3 
-      dec:px-4 
-      dec:py-3;
+      dec:px-3 
+      dec:py-2;
   }
 
   .dec-root .dec-task-node-icon {
+    @apply dec:shrink-0;
     color: var(--task-node-color);
   }
 
   .dec-root .dec-task-node-label {
     @apply dec:flex 
     dec:flex-col 
-    dec:gap-0.5;
+    dec:gap-0.5
+    dec:min-w-0;
   }
 
   .dec-root .dec-task-node-name {
-    @apply dec:text-sm 
+    @apply dec:truncate /*TODO: for now truncate text revisit when layout is in and working on styling tweaks */
+    dec:text-sm 
     dec:text-black 
     dec:leading-tight;
   }
@@ -232,4 +235,28 @@
   .dec-root.dark .dec-task-node-type {
     @apply dec:text-gray-400;
   }
+
+  .dec-root .dec-task-node-badge {
+    @apply dec:ml-auto
+      dec:shrink-0 
+      dec:rounded
+      dec:px-2 
+      dec:py-0.5 
+      dec:text-[8px]
+      dec:font-semibold
+      dec:uppercase
+      dec:whitespace-nowrap;
+    color: var(--task-node-color);
+    border: 1px solid var(--task-node-color);
+  }
+
+  .dec-root .dec-task-node-badge-icon {
+    @apply dec:ml-auto
+      dec:shrink-0 
+      dec:flex 
+      dec:items-center 
+      dec:justify-center;
+    color: var(--task-node-color);
+  }
+  /* end task leaf nodes */
 }
diff --git a/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/autoLayout.ts b/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/autoLayout.ts
index 4f96b2d..7af17a3 100644
--- a/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/autoLayout.ts
+++ b/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/autoLayout.ts
@@ -19,7 +19,7 @@ import { ReactFlowGraph } from "./diagramBuilder";
 // Defaults
 export const DEFAULT_NODE_SIZE = {
   height: 60,
-  width: 180,
+  width: 200,
 };
 
 export type Point = {
diff --git a/packages/serverless-workflow-diagram-editor/src/react-flow/nodes/Nodes.tsx b/packages/serverless-workflow-diagram-editor/src/react-flow/nodes/Nodes.tsx
index aca8f10..f202ab8 100644
--- a/packages/serverless-workflow-diagram-editor/src/react-flow/nodes/Nodes.tsx
+++ b/packages/serverless-workflow-diagram-editor/src/react-flow/nodes/Nodes.tsx
@@ -18,6 +18,8 @@ import type React from "react";
 import { GraphNodeType, type Specification } from "@serverlessworkflow/sdk";
 import * as RF from "@xyflow/react";
 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 ReactFlowNodeTypes: RF.NodeTypes = {
@@ -33,10 +35,28 @@ export const ReactFlowNodeTypes: RF.NodeTypes = {
   [GraphNodeType.Run]: RunNode,
   [GraphNodeType.Set]: SetNode,
   [GraphNodeType.Switch]: SwitchNode,
+  [GraphNodeType.TryCatch]: TryCatchNode,
   [GraphNodeType.Try]: TryNode,
+  [GraphNodeType.Catch]: CatchNode,
   [GraphNodeType.Wait]: WaitNode,
 };
 
+const KNOWN_BADGES = new Set([
+  "http",
+  "grpc",
+  "asyncapi",
+  "openapi",
+  "a2a",
+  "mcp",
+  "container",
+  "script",
+  "shell",
+  "workflow",
+  "all",
+  "any",
+  "one",
+]);
+
 export type BaseNodeData<T = Specification.Task | void> = {
   label: string;
   task?: T;
@@ -47,9 +67,34 @@ interface NodeContentProps {
   data: BaseNodeData;
   selected: boolean;
   type: string;
+  badge?: string | undefined;
+}
+
+interface BadgeProps {
+  badge: string;
+  testId: string;
+}
+
+function TaskNodeBadge({ badge, testId }: BadgeProps) {
+  const isUnknown = !KNOWN_BADGES.has(badge.toLowerCase());
+
+  if (isUnknown) {
+    /* TODO: instead of using the browser default to display tool tip like below, replace with tooltip component when we add it */
+    return (
+      <span title={badge} className="dec-task-node-badge-icon" data-testid={`${testId}-icon`}>
+        <Info size={18} />
+      </span>
+    );
+  }
+
+  return (
+    <span className="dec-task-node-badge" data-testid={testId}>
+      {badge}
+    </span>
+  );
 }
 
-function TaskNodeContent({ id, data, selected, type }: NodeContentProps) {
+function TaskNodeContent({ id, data, selected, type, badge }: NodeContentProps) {
   const config = taskNodeConfigMap[type as LeafNodeType];
   const Icon = config.icon;
   return (
@@ -65,6 +110,7 @@ function TaskNodeContent({ id, data, selected, type }: NodeContentProps) {
           <span className="dec-task-node-name">{data.label}</span>
           <span className="dec-task-node-type">{config.typeLabel}</span>
         </div>
+        {badge && <TaskNodeBadge badge={badge} testId={`${type}-node-${id}-badge`} />}
       </div>
       <RF.Handle type="source" position={RF.Position.Bottom} />
     </div>
@@ -109,49 +155,51 @@ export function EndNode({ id, data, selected, type }: RF.NodeProps<EndNodeType>)
   return <PlaceholderContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* call node */
+/* call leaf node */
 export type CallNodeType = RF.Node<BaseNodeData<Specification.CallTask>, typeof GraphNodeType.Call>;
 export function CallNode({ id, data, selected, type }: RF.NodeProps<CallNodeType>) {
-  return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
+  const badge = data.task ? getCallSubType(data.task) : undefined;
+  return <TaskNodeContent id={id} data={data} selected={selected} type={type} badge={badge} />;
 }
 
-/* do node */
+/* do container node */
 export type DoNodeType = RF.Node<BaseNodeData<Specification.DoTask>, typeof GraphNodeType.Do>;
 export function DoNode({ id, data, selected, type }: RF.NodeProps<DoNodeType>) {
   // TODO: This component is just a placeholder
   return <PlaceholderContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* emit node */
+/* emit leaf node */
 export type EmitNodeType = RF.Node<BaseNodeData<Specification.EmitTask>, typeof GraphNodeType.Emit>;
 export function EmitNode({ id, data, selected, type }: RF.NodeProps<EmitNodeType>) {
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* for node */
+/* for container node */
 export type ForNodeType = RF.Node<BaseNodeData<Specification.ForTask>, typeof GraphNodeType.For>;
 export function ForNode({ id, data, selected, type }: RF.NodeProps<ForNodeType>) {
   // TODO: This component is just a placeholder
   return <PlaceholderContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* fork node */
+/* fork container node */
 export type ForkNodeType = RF.Node<BaseNodeData<Specification.ForkTask>, typeof GraphNodeType.Fork>;
 export function ForkNode({ id, data, selected, type }: RF.NodeProps<ForkNodeType>) {
   // TODO: This component is just a placeholder
   return <PlaceholderContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* listen node */
+/* listen leaf node */
 export type ListenNodeType = RF.Node<
   BaseNodeData<Specification.ListenTask>,
   typeof GraphNodeType.Listen
 >;
 export function ListenNode({ id, data, selected, type }: RF.NodeProps<ListenNodeType>) {
-  return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
+  const badge = data.task ? getListenSubType(data.task) : undefined;
+  return <TaskNodeContent id={id} data={data} selected={selected} type={type} badge={badge} />;
 }
 
-/* raise node */
+/* raise leaf node */
 export type RaiseNodeType = RF.Node<
   BaseNodeData<Specification.RaiseTask>,
   typeof GraphNodeType.Raise
@@ -160,19 +208,20 @@ export function RaiseNode({ id, data, selected, type }: RF.NodeProps<RaiseNodeTy
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* run node */
+/* run leaf node */
 export type RunNodeType = RF.Node<BaseNodeData<Specification.RunTask>, typeof GraphNodeType.Run>;
 export function RunNode({ id, data, selected, type }: RF.NodeProps<RunNodeType>) {
-  return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
+  const badge = data.task ? getRunSubType(data.task) : undefined;
+  return <TaskNodeContent id={id} data={data} selected={selected} type={type} badge={badge} />;
 }
 
-/* set node */
+/* set leaf node */
 export type SetNodeType = RF.Node<BaseNodeData<Specification.SetTask>, typeof GraphNodeType.Set>;
 export function SetNode({ id, data, selected, type }: RF.NodeProps<SetNodeType>) {
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* switch node */
+/* switch leaf node */
 export type SwitchNodeType = RF.Node<
   BaseNodeData<Specification.SwitchTask>,
   typeof GraphNodeType.Switch
@@ -181,14 +230,27 @@ export function SwitchNode({ id, data, selected, type }: RF.NodeProps<SwitchNode
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* try node */
+/* try catch container node */
+export type TryCatchNodeType = RF.Node<BaseNodeData, typeof GraphNodeType.TryCatch>;
+export function TryCatchNode({ id, data, selected, type }: RF.NodeProps<TryCatchNodeType>) {
+  // TODO: This component is just a placeholder
+  return <PlaceholderContent id={id} data={data} selected={selected} type={type} />;
+}
+
+/* try container node */
 export type TryNodeType = RF.Node<BaseNodeData<Specification.TryTask>, typeof GraphNodeType.Try>;
 export function TryNode({ id, data, selected, type }: RF.NodeProps<TryNodeType>) {
   // TODO: This component is just a placeholder
   return <PlaceholderContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* wait node */
+/* catch leaf node */
+export type CatchNodeType = RF.Node<BaseNodeData, typeof GraphNodeType.Catch>;
+export function CatchNode({ id, data, selected, type }: RF.NodeProps<CatchNodeType>) {
+  return <TaskNodeContent 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>) {
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
diff --git a/packages/serverless-workflow-diagram-editor/src/react-flow/nodes/taskNodeConfig.ts b/packages/serverless-workflow-diagram-editor/src/react-flow/nodes/taskNodeConfig.ts
index b71a715..5580651 100644
--- a/packages/serverless-workflow-diagram-editor/src/react-flow/nodes/taskNodeConfig.ts
+++ b/packages/serverless-workflow-diagram-editor/src/react-flow/nodes/taskNodeConfig.ts
@@ -23,12 +23,14 @@ import {
   Megaphone,
   PenLine,
   Phone,
+  ShieldAlert,
   Terminal,
 } from "lucide-react";
 import type { ComponentType } from "react";
 
 export type LeafNodeType =
   | typeof GraphNodeType.Call
+  | typeof GraphNodeType.Catch
   | typeof GraphNodeType.Emit
   | typeof GraphNodeType.Listen
   | typeof GraphNodeType.Raise
@@ -49,6 +51,11 @@ export const taskNodeConfigMap: Record<LeafNodeType, TaskNodeConfig> = {
     icon: Phone,
     typeLabel: "CALL",
   },
+  [GraphNodeType.Catch]: {
+    color: "#F97316",
+    icon: ShieldAlert,
+    typeLabel: "CATCH",
+  },
   [GraphNodeType.Emit]: {
     color: "#8B5CF6",
     icon: Megaphone,
diff --git a/packages/serverless-workflow-diagram-editor/stories/DiagramEditor.stories.ts b/packages/serverless-workflow-diagram-editor/stories/DiagramEditor.stories.ts
index 44f8d77..1301837 100644
--- a/packages/serverless-workflow-diagram-editor/stories/DiagramEditor.stories.ts
+++ b/packages/serverless-workflow-diagram-editor/stories/DiagramEditor.stories.ts
@@ -17,7 +17,64 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
 
 import { DiagramEditor } from "./DiagramEditor";
-import { BASIC_VALID_WORKFLOW_YAML } from "../tests/fixtures/workflows";
+
+const workflowExample = `document:
+  dsl: '1.0.3'
+  namespace: examples
+  name: accumulate-room-readings
+  version: '0.1.0'
+do:
+  - consumeReading:
+      listen:
+        to:
+          all:
+            - with:
+                source: https://my.home.com/sensor
+                type: my.home.sensors.temperature
+              correlate:
+                roomId:
+                  from: .roomid
+            - with:
+                source: https://my.home.com/sensor
+                type: my.home.sensors.humidity
+              correlate:
+                roomId:
+                  from: .roomid
+      output:
+        as: .data.reading
+  - logReading:
+      for:
+        each: reading
+        in: .readings
+      do:
+        - callOrderService:
+            call: openapi
+            with:
+              document:
+                endpoint: http://myorg.io/ordersservices.json
+              operationId: logreading
+  - generateReport:
+      call: openapi
+      with:
+        document:
+          endpoint: http://myorg.io/ordersservices.json
+        operationId: produceReport
+  - emitEvent:
+      emit:
+        event:
+          with:
+            source: https://petstore.com
+            type: com.petstore.order.placed.v1
+            data:
+              client:
+                firstName: Cruella
+                lastName: de Vil
+              items:
+                - breed: dalmatian
+                  quantity: 101
+timeout:
+  after:
+    hours: 1`;
 
 const meta = {
   id: "diagram-editor",
@@ -40,6 +97,6 @@ export const Component: Story = {
     isReadOnly: true,
     locale: "en",
     colorMode: "system",
-    content: BASIC_VALID_WORKFLOW_YAML, // TODO: Add better workflow sample when removing hardcoded nodes and edges in Diagram component
+    content: workflowExample,
   },
 };
diff --git a/packages/serverless-workflow-diagram-editor/tests-e2e/diagram-editor.spec.ts b/packages/serverless-workflow-diagram-editor/tests-e2e/diagram-editor.spec.ts
index 85f7612..371d080 100644
--- a/packages/serverless-workflow-diagram-editor/tests-e2e/diagram-editor.spec.ts
+++ b/packages/serverless-workflow-diagram-editor/tests-e2e/diagram-editor.spec.ts
@@ -23,13 +23,13 @@ test("diagram editor renders correctly", async ({ page }) => {
   await expect(page.getByTestId("diagram-container")).toBeVisible();
 
   // Check at least one specific node
-  await expect(page.getByTestId("rf__node-/do/0/step1")).toContainText("step1SET");
+  await expect(page.getByTestId("rf__node-/do/0/consumeReading")).toContainText("consumeReading");
 
   // Check total nodes
   const nodes = page.locator('[data-testid^="rf__node-"]');
-  await expect(nodes).toHaveCount(3);
+  await expect(nodes).toHaveCount(6);
 
   // Check total edge
   const edges = page.locator('[data-testid^="rf__edge-"]');
-  await expect(edges).toHaveCount(2);
+  await expect(edges).toHaveCount(5);
 });
diff --git a/packages/serverless-workflow-diagram-editor/tests/core/taskSubType.test.ts b/packages/serverless-workflow-diagram-editor/tests/core/taskSubType.test.ts
new file mode 100644
index 0000000..4a1b852
--- /dev/null
+++ b/packages/serverless-workflow-diagram-editor/tests/core/taskSubType.test.ts
@@ -0,0 +1,79 @@
+/*
+ * 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 type { Specification } from "@serverlessworkflow/sdk";
+import { describe, expect, it } from "vitest";
+import { getCallSubType, getListenSubType, getRunSubType } from "../../src/core";
+
+describe("getCallSubType", () => {
+  it.each([
+    ["http", "http"],
+    ["grpc", "grpc"],
+    ["asyncapi", "asyncapi"],
+    ["openapi", "openapi"],
+    ["a2a", "a2a"],
+    ["mcp", "mcp"],
+    ["myCustomFunction", "myCustomFunction"],
+  ])("should return '%s' for call subtype '%s'", (callValue, expectedSubType) => {
+    const task = { call: callValue } as Specification.CallTask;
+    expect(getCallSubType(task)).toBe(expectedSubType);
+  });
+
+  it("should return undefined when call is missing", () => {
+    expect(getCallSubType({} as Specification.CallTask)).toBeUndefined();
+  });
+});
+
+describe("getRunSubType", () => {
+  it.each([
+    ["container", "container"],
+    ["script", "script"],
+    ["shell", "shell"],
+    ["workflow", "workflow"],
+    ["unknownRunType", "unknownRunType"],
+  ])("should return '%s' for run subtype '%s'", (runKey, expectedSubType) => {
+    const task = { run: { [runKey]: {} } } as unknown as Specification.RunTask;
+    expect(getRunSubType(task)).toBe(expectedSubType);
+  });
+
+  it.each([
+    ["run is missing", {}],
+    ["run is not an object", { run: "invalidRunValue" }],
+    ["run is an array", { run: [] }],
+  ])("should return undefined when %s", (_label, task) => {
+    expect(getRunSubType(task as unknown as Specification.RunTask)).toBeUndefined();
+  });
+});
+
+describe("getListenSubType", () => {
+  it.each([
+    ["all", "all"],
+    ["any", "any"],
+    ["one", "one"],
+    ["unknownListenType", "unknownListenType"],
+  ])("should return '%s' for listen subtype '%s'", (listenKey, expectedSubType) => {
+    const task = { listen: { to: { [listenKey]: {} } } } as unknown as Specification.ListenTask;
+    expect(getListenSubType(task)).toBe(expectedSubType);
+  });
+
+  it.each([
+    ["listen is missing", {}],
+    ["listen.to is missing", { listen: {} }],
+    ["listen.to is an array", { listen: { to: [] } }],
+  ])("should return undefined when %s", (_label, task) => {
+    expect(getListenSubType(task as unknown as Specification.ListenTask)).toBeUndefined();
+  });
+});
diff --git a/packages/serverless-workflow-diagram-editor/tests/react-flow/nodes/Nodes.test.tsx b/packages/serverless-workflow-diagram-editor/tests/react-flow/nodes/Nodes.test.tsx
index 1245987..26a93c4 100644
--- a/packages/serverless-workflow-diagram-editor/tests/react-flow/nodes/Nodes.test.tsx
+++ b/packages/serverless-workflow-diagram-editor/tests/react-flow/nodes/Nodes.test.tsx
@@ -22,14 +22,20 @@ import { ReactFlowNodeTypes } from "../../../src/react-flow/nodes/Nodes";
 import { taskNodeConfigMap, type LeafNodeType } from "../../../src/react-flow/nodes/taskNodeConfig";
 import { DEFAULT_NODE_SIZE } from "../../../src/react-flow/diagram/autoLayout";
 
-function testNode(id: string, type: string, y: number, label: string): RF.Node {
+function testNode(
+  id: string,
+  type: string,
+  y: number,
+  label: string,
+  task?: Record<string, unknown>,
+): RF.Node {
   return {
     id,
     type,
     position: { x: 100, y },
     height: DEFAULT_NODE_SIZE.height,
     width: DEFAULT_NODE_SIZE.width,
-    data: { label },
+    data: { label, ...(task !== undefined ? { task } : {}) },
   };
 }
 
@@ -45,8 +51,10 @@ const allNodes: RF.Node[] = [
   testNode("n8", GraphNodeType.Raise, 700, "Node 8"),
   testNode("n9", GraphNodeType.Run, 800, "Node 9"),
   testNode("n10", GraphNodeType.Set, 900, "Node 10"),
-  testNode("n11", GraphNodeType.Try, 1000, "Node 11"),
-  testNode("n12", GraphNodeType.Wait, 1100, "Node 12"),
+  testNode("n11", GraphNodeType.TryCatch, 1000, "Node 11"),
+  testNode("n12", GraphNodeType.Try, 1100, "Node 12"),
+  testNode("n13", GraphNodeType.Catch, 1200, "Node 13"),
+  testNode("n14", GraphNodeType.Wait, 1300, "Node 14"),
   testNode("end", GraphNodeType.End, 0, "End"),
 ];
 
@@ -55,17 +63,17 @@ const allEdges: RF.Edge[] = [
   { id: "n1-n2", source: "n1", target: "n2" },
   { id: "n2-n3", source: "n2", target: "n3" },
   { id: "n3-n4", source: "n3", target: "n4" },
-  { id: "n3-n5", source: "n3", target: "n5" },
-  { id: "n3-n6", source: "n3", target: "n6" },
-  { id: "n4-n7", source: "n4", target: "n7" },
-  { id: "n5-n7", source: "n5", target: "n7" },
+  { id: "n4-n5", source: "n4", target: "n5" },
+  { id: "n5-n6", source: "n5", target: "n6" },
   { id: "n6-n7", source: "n6", target: "n7" },
   { id: "n7-n8", source: "n7", target: "n8" },
   { id: "n8-n9", source: "n8", target: "n9" },
   { id: "n9-n10", source: "n9", target: "n10" },
   { id: "n10-n11", source: "n10", target: "n11" },
   { id: "n11-n12", source: "n11", target: "n12" },
-  { id: "n12-end", source: "n12", target: "end" },
+  { id: "n12-n13", source: "n12", target: "n13" },
+  { id: "n13-n14", source: "n13", target: "n14" },
+  { id: "n14-end", source: "n14", target: "end" },
 ];
 
 describe("React Flow custom node types", () => {
@@ -91,8 +99,10 @@ describe("React Flow custom node types", () => {
     expect(screen.getByTestId("raise-node-n8")).toBeInTheDocument();
     expect(screen.getByTestId("run-node-n9")).toBeInTheDocument();
     expect(screen.getByTestId("set-node-n10")).toBeInTheDocument();
-    expect(screen.getByTestId("try-node-n11")).toBeInTheDocument();
-    expect(screen.getByTestId("wait-node-n12")).toBeInTheDocument();
+    expect(screen.getByTestId("try-catch-node-n11")).toBeInTheDocument();
+    expect(screen.getByTestId("try-node-n12")).toBeInTheDocument();
+    expect(screen.getByTestId("catch-node-n13")).toBeInTheDocument();
+    expect(screen.getByTestId("wait-node-n14")).toBeInTheDocument();
     expect(screen.getByTestId("end-node-end")).toBeInTheDocument();
   });
 
@@ -105,7 +115,8 @@ describe("React Flow custom node types", () => {
       { id: "n8", type: GraphNodeType.Raise, testId: "raise" },
       { id: "n9", type: GraphNodeType.Run, testId: "run" },
       { id: "n10", type: GraphNodeType.Set, testId: "set" },
-      { id: "n12", type: GraphNodeType.Wait, testId: "wait" },
+      { id: "n13", type: GraphNodeType.Catch, testId: "catch" },
+      { id: "n14", type: GraphNodeType.Wait, testId: "wait" },
     ];
 
     it.each(leafNodes)("should render %s node with correct config", ({ id, type, testId }) => {
@@ -125,4 +136,62 @@ describe("React Flow custom node types", () => {
       expect(node.style.getPropertyValue("--task-node-color")).toBe(config.color);
     });
   });
+
+  describe("badge rendering", () => {
+    it("should render a text badge for known subtypes", () => {
+      const nodesWithBadges = [
+        testNode("n1", GraphNodeType.Call, 10, "CallNode", { call: "http" }),
+        testNode("n2", GraphNodeType.Listen, 100, "ListenNode", { listen: { to: { any: [] } } }),
+      ];
+      render(
+        <div>
+          <RF.ReactFlow nodeTypes={ReactFlowNodeTypes} nodes={nodesWithBadges} edges={allEdges} />
+        </div>,
+      );
+
+      const callBadge = screen.getByTestId("call-node-n1-badge");
+      expect(callBadge).toBeInTheDocument();
+      expect(callBadge.textContent).toBe("http");
+
+      const listenBadge = screen.getByTestId("listen-node-n2-badge");
+      expect(listenBadge).toBeInTheDocument();
+      expect(listenBadge.textContent).toBe("any");
+    });
+
+    it("should render an icon badge for an unknown subtype", () => {
+      const nodesWithUnknownBadges = [
+        testNode("n1", GraphNodeType.Call, 100, "CallNode", { call: "customCall" }),
+      ];
+      render(
+        <div>
+          <RF.ReactFlow
+            nodeTypes={ReactFlowNodeTypes}
+            nodes={nodesWithUnknownBadges}
+            edges={allEdges}
+          />
+        </div>,
+      );
+
+      const callBadge = screen.getByTestId("call-node-n1-badge-icon");
+      expect(callBadge).toBeInTheDocument();
+    });
+
+    it("should not render a badge when task has no subtype", () => {
+      const nodesWithoutBadges = [
+        testNode("n1", GraphNodeType.Wait, 100, "WaitNode", { wait: "PT1S" }),
+      ];
+      render(
+        <div>
+          <RF.ReactFlow
+            nodeTypes={ReactFlowNodeTypes}
+            nodes={nodesWithoutBadges}
+            edges={allEdges}
+          />
+        </div>,
+      );
+
+      const badge = screen.queryByTestId("wait-node-n1-badge");
+      expect(badge).not.toBeInTheDocument();
+    });
+  });
 });
diff --git a/packages/serverless-workflow-diagram-editor/tests/react-flow/nodes/taskNodeConfig.test.ts b/packages/serverless-workflow-diagram-editor/tests/react-flow/nodes/taskNodeConfig.test.ts
index 19ff1e4..6f0b6b9 100644
--- a/packages/serverless-workflow-diagram-editor/tests/react-flow/nodes/taskNodeConfig.test.ts
+++ b/packages/serverless-workflow-diagram-editor/tests/react-flow/nodes/taskNodeConfig.test.ts
@@ -27,6 +27,7 @@ const leafNodeTypes: LeafNodeType[] = [
   GraphNodeType.Set,
   GraphNodeType.Switch,
   GraphNodeType.Wait,
+  GraphNodeType.Catch,
 ];
 
 describe("taskNodeConfig", () => {