diff --git a/packages/serverless-workflow-diagram-editor/src/core/autoLayout.ts b/packages/serverless-workflow-diagram-editor/src/core/autoLayout.ts
index b30de77..e270e5c 100644
--- a/packages/serverless-workflow-diagram-editor/src/core/autoLayout.ts
+++ b/packages/serverless-workflow-diagram-editor/src/core/autoLayout.ts
@@ -19,7 +19,7 @@ import { ExtendedGraph, Position, Size } from "./graph";
 // Defaults
 export const DEFAULT_NODE_SIZE = {
   height: 60,
-  width: 180,
+  width: 200,
 };
 
 export function applyAutoLayout(graph: ExtendedGraph): ExtendedGraph {
diff --git a/packages/serverless-workflow-diagram-editor/src/core/index.ts b/packages/serverless-workflow-diagram-editor/src/core/index.ts
index 800557d..c1be979 100644
--- a/packages/serverless-workflow-diagram-editor/src/core/index.ts
+++ b/packages/serverless-workflow-diagram-editor/src/core/index.ts
@@ -17,3 +17,4 @@
 export * from "./workflowSdk";
 export * from "./graph";
 export * from "./autoLayout";
+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..aae4736
--- /dev/null
+++ b/packages/serverless-workflow-diagram-editor/src/core/taskSubType.ts
@@ -0,0 +1,53 @@
+/*
+ * 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 { GraphNodeType, type Specification } from "@serverlessworkflow/sdk";
+
+export function getCallSubType(task: Specification.CallTask): string | undefined {
+  return typeof task.call === "string" ? task.call : undefined;
+}
+
+export function getRunSubType(task: Specification.RunTask): string | undefined {
+  const run = task.run;
+  if (run && typeof run === "object") {
+    const firstKey = Object.keys(run)[0];
+    return firstKey ?? undefined;
+  }
+  return undefined;
+}
+
+export function getListenSubType(task: Specification.ListenTask): string | undefined {
+  const listen = task.listen?.to;
+  if (listen && typeof listen === "object") {
+    const firstKey = Object.keys(listen)[0];
+    return firstKey ?? undefined;
+  }
+  return undefined;
+}
+
+/* TODO: Add container subtypes when container nodes are available. This is the entry point to be called when we remove hardcoded values in Diagram.tsx */
+export function getTaskSubType(nodeType: GraphNodeType, task: Specification.Task): string | undefined {
+  switch (nodeType) {
+    case GraphNodeType.Call:
+      return getCallSubType(task as Specification.CallTask);
+    case GraphNodeType.Run:
+      return getRunSubType(task as Specification.RunTask);
+    case GraphNodeType.Listen:
+      return getListenSubType(task as Specification.ListenTask);
+    default:
+      return 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..c1b0011 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
+    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/Diagram.tsx b/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/Diagram.tsx
index 64a24ce..ce6a71f 100644
--- a/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/Diagram.tsx
+++ b/packages/serverless-workflow-diagram-editor/src/react-flow/diagram/Diagram.tsx
@@ -39,7 +39,7 @@ const initialNodes: RF.Node[] = [
     position: { x: 100, y: 0 },
     height: DEFAULT_NODE_SIZE.height,
     width: DEFAULT_NODE_SIZE.width,
-    data: { label: "CallNode" },
+    data: { label: "CallNode", badge: "HTTP" },
   },
   {
     id: "n2",
@@ -60,7 +60,7 @@ const initialNodes: RF.Node[] = [
   {
     id: "n4",
     type: GraphNodeType.Emit,
-    position: { x: -100, y: 300 },
+    position: { x: -150, y: 300 },
     height: DEFAULT_NODE_SIZE.height,
     width: DEFAULT_NODE_SIZE.width,
     data: { label: "EmitNode" },
@@ -76,7 +76,7 @@ const initialNodes: RF.Node[] = [
   {
     id: "n6",
     type: GraphNodeType.Fork,
-    position: { x: 300, y: 300 },
+    position: { x: 350, y: 300 },
     height: DEFAULT_NODE_SIZE.height,
     width: DEFAULT_NODE_SIZE.width,
     data: { label: "Node 6" },
@@ -87,7 +87,7 @@ const initialNodes: RF.Node[] = [
     position: { x: 100, y: 400 },
     height: DEFAULT_NODE_SIZE.height,
     width: DEFAULT_NODE_SIZE.width,
-    data: { label: "ListenNode" },
+    data: { label: "ListenNode", badge: "ALL" },
   },
   {
     id: "n8",
@@ -103,7 +103,7 @@ const initialNodes: RF.Node[] = [
     position: { x: 100, y: 600 },
     height: DEFAULT_NODE_SIZE.height,
     width: DEFAULT_NODE_SIZE.width,
-    data: { label: "RunNode" },
+    data: { label: "RunNode", badge: "myCustomType" },
   },
   {
     id: "n10",
@@ -139,12 +139,11 @@ const initialEdges: RF.Edge[] = [
     type: GraphEdgeType.Default,
     data: {
       wayPoints: [
-        { x: 190, y: 60 },
-        { x: 190, y: 70 },
+        { x: 200, y: 70 }, 
         { x: 140, y: 70 },
         { x: 140, y: 85 },
-        { x: 190, y: 85 },
-        { x: 190, y: 95 }, 
+        { x: 200, y: 85 },
+        { x: 200, y: 95 },
       ],
     },
   },
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 d5f7bb0..9882560 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,7 @@ import type React from "react";
 import { GraphNodeType } from "@serverlessworkflow/sdk";
 import * as RF from "@xyflow/react";
 import { type LeafNodeType, taskNodeConfigMap } from "./taskNodeConfig";
+import { Info } from "lucide-react";
 
 // Node types must match sdk GraphNodeType enum
 export const NodeTypes: RF.NodeTypes = {
@@ -31,13 +32,32 @@ export const NodeTypes: 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 = {
   // TODO: It is a placeholder, add properties to be consumed internally by node components
   label: string;
+  badge?: string;
 };
 
 interface NodeContentProps {
@@ -47,6 +67,30 @@ interface NodeContentProps {
   type: string;
 }
 
+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) {
   const config = taskNodeConfigMap[type as LeafNodeType];
   const Icon = config.icon;
@@ -63,6 +107,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>
+        {data.badge && <TaskNodeBadge badge={data.badge} testId={`${type}-node-${id}-badge`} />}
       </div>
       <RF.Handle type="source" position={RF.Position.Bottom} />
     </div>
@@ -93,77 +138,90 @@ function PlaceholderContent({ id, data, selected, type }: PlaceholderProps) {
   );
 }
 
-/* call node */
+/* call leaf node */
 export type CallNodeType = RF.Node<BaseNodeData, typeof GraphNodeType.Call>;
 export function CallNode({ id, data, selected, type }: RF.NodeProps<CallNodeType>) {
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* do node */
+/* do container node */
 export type DoNodeType = RF.Node<BaseNodeData, 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, 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, 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, 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, typeof GraphNodeType.Listen>;
 export function ListenNode({ id, data, selected, type }: RF.NodeProps<ListenNodeType>) {
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* raise node */
+/* raise leaf node */
 export type RaiseNodeType = RF.Node<BaseNodeData, typeof GraphNodeType.Raise>;
 export function RaiseNode({ id, data, selected, type }: RF.NodeProps<RaiseNodeType>) {
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* run node */
+/* run leaf node */
 export type RunNodeType = RF.Node<BaseNodeData, typeof GraphNodeType.Run>;
 export function RunNode({ id, data, selected, type }: RF.NodeProps<RunNodeType>) {
   return <TaskNodeContent id={id} data={data} selected={selected} type={type} />;
 }
 
-/* set node */
+/* set leaf node */
 export type SetNodeType = RF.Node<BaseNodeData, 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, typeof GraphNodeType.Switch>;
 export function SwitchNode({ id, data, selected, type }: RF.NodeProps<SwitchNodeType>) {
   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, 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, 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/tests/core/taskSubType.test.ts b/packages/serverless-workflow-diagram-editor/tests/core/taskSubType.test.ts
new file mode 100644
index 0000000..76aeb5e
--- /dev/null
+++ b/packages/serverless-workflow-diagram-editor/tests/core/taskSubType.test.ts
@@ -0,0 +1,124 @@
+/*
+ * 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 { GraphNodeType, type Specification } from "@serverlessworkflow/sdk";
+import { describe, expect, it } from "vitest";
+import { getCallSubType, getListenSubType, getRunSubType, getTaskSubType } from "../../src/core";
+
+describe("getTaskSubType", () => {
+  it("should call getCallSubType for CallTask", () => {
+    const task = { call: "http" } as Specification.CallTask;
+    expect(getTaskSubType(GraphNodeType.Call, task)).toBe("http");
+  });
+
+  it("should call getRunSubType for RunTask", () => {
+    const task = { run: { script: {} } } as Specification.RunTask;
+    expect(getTaskSubType(GraphNodeType.Run, task)).toBe("script");
+  });
+
+  it("should call getListenSubType for ListenTask", () => {
+    const task = { listen: { to: { any: {} } } } as Specification.ListenTask;
+    expect(getTaskSubType(GraphNodeType.Listen, task)).toBe("any");
+  });
+
+  it("should return undefined for nodes without subtypes", () => {
+    expect(getTaskSubType(GraphNodeType.Emit, {} as Specification.Task)).toBeUndefined();
+    expect(getTaskSubType(GraphNodeType.Raise, {} as Specification.Task)).toBeUndefined();
+    expect(getTaskSubType(GraphNodeType.Set, {} as Specification.Task)).toBeUndefined();
+    expect(getTaskSubType(GraphNodeType.Switch, {} as Specification.Task)).toBeUndefined();
+    expect(getTaskSubType(GraphNodeType.Try, {} as Specification.Task)).toBeUndefined();
+    expect(getTaskSubType(GraphNodeType.Wait, {} as Specification.Task)).toBeUndefined();
+  });
+});
+
+describe("getCallSubType", () => {
+  it.each([
+    ["http", "http"],
+    ["grpc", "grpc"],
+    ["asyncapi", "asyncapi"],
+    ["openapi", "openapi"],
+    ["a2a", "a2a"],
+    ["mcp", "mcp"],
+  ])("should return %s correct for known call subtype '%s':", (callValue, expectedSubType) => {
+    const task = { call: callValue } as Specification.CallTask;
+    expect(getCallSubType(task)).toBe(expectedSubType);
+  });
+
+  it("should return the custom function name for unknown subtypes", () => {
+    const task = { call: "myCustomFunction" } as Specification.CallTask;
+    expect(getCallSubType(task)).toBe("myCustomFunction");
+  });
+
+  it("should return undefined when call is missing", () => {
+    const task = {} as Specification.CallTask;
+    expect(getCallSubType(task)).toBeUndefined();
+  });
+});
+
+describe("getRunSubType", () => {
+  it.each([
+    ["container", "container"],
+    ["script", "script"],
+    ["shell", "shell"],
+    ["workflow", "workflow"],
+  ])("should return '%s' run key subtype  %s", (runKey, expectedSubType) => {
+    const task = { run: { [runKey]: {} } } as Specification.RunTask;
+    expect(getRunSubType(task)).toBe(expectedSubType);
+  });
+
+  it("should return unknown run types", () => {
+    const task = { run: { unknownRunType: {} } } as unknown as Specification.RunTask;
+    expect(getRunSubType(task)).toBe("unknownRunType");
+  });
+
+  it("should return undefined when run is missing", () => {
+    const task = {} as Specification.RunTask;
+    expect(getRunSubType(task)).toBeUndefined();
+  });
+
+  it("should return undefined when run is not an object", () => {
+    const task = { run: "invalidRunValue" } as unknown as Specification.RunTask;
+    expect(getRunSubType(task)).toBeUndefined();
+  });
+});
+
+describe("getListenSubType", () => {
+  it.each([
+    ["all", "all"],
+    ["any", "any"],
+    ["one", "one"],
+  ])("should return '%s' listen key subtype %s", (listenKey, expectedSubType) => {
+    const task = { listen: { to: { [listenKey]: {} } } } as Specification.ListenTask;
+    expect(getListenSubType(task)).toBe(expectedSubType);
+  });
+
+  it("should return unknown listen types", () => {
+    const task = {
+      listen: { to: { unknownListenType: {} } },
+    } as unknown as Specification.ListenTask;
+    expect(getListenSubType(task)).toBe("unknownListenType");
+  });
+
+  it("should return undefined when listen is missing", () => {
+    const task = {} as Specification.ListenTask;
+    expect(getListenSubType(task)).toBeUndefined();
+  });
+
+  it("should return undefined when listen.to is missing", () => {
+    const task = { listen: {} } as Specification.ListenTask;
+    expect(getListenSubType(task)).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 b5a1896..7af8412 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,14 @@ import { NodeTypes } from "../../../src/react-flow/nodes/Nodes";
 import { DEFAULT_NODE_SIZE } from "../../../src/core";
 import { taskNodeConfigMap, type LeafNodeType } from "../../../src/react-flow/nodes/taskNodeConfig";
 
-function testNode(id: string, type: string, y: number, label: string): RF.Node {
+function testNode(id: string, type: string, y: number, label: string, badge?: string): RF.Node {
   return {
     id,
     type,
     position: { x: 100, y },
     height: DEFAULT_NODE_SIZE.height,
     width: DEFAULT_NODE_SIZE.width,
-    data: { label },
+    data: { label, ...(badge !== undefined ? { badge } : {}) },
   };
 }
 
@@ -44,24 +44,26 @@ 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"),
 ];
 
 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-n13", source: "n12", target: "n13" },
+  { id: "n13-n14", source: "n13", target: "n14" },
 ];
 
 describe("React Flow custom node types", () => {
@@ -86,8 +88,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();
   });
 
   describe("should render leaf nodes with TaskNodeContent", () => {
@@ -99,7 +103,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 }) => {
@@ -119,4 +124,56 @@ 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", "HTTP"),
+        testNode("n2", GraphNodeType.Listen, 100, "ListenNode", "ANY"),
+      ];
+      render(
+        <div>
+          <RF.ReactFlow nodeTypes={NodeTypes} 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", "UnknownCall"),
+        testNode("n2", GraphNodeType.Listen, 100, "ListenNode", "UnknownListen"),
+      ];
+      render(
+        <div>
+          <RF.ReactFlow nodeTypes={NodeTypes} nodes={nodesWithUnknownBadges} edges={allEdges} />
+        </div>,
+      );
+
+      const callBadge = screen.getByTestId("call-node-n1-badge-icon");
+      expect(callBadge).toBeInTheDocument();
+
+      const listenBadge = screen.getByTestId("listen-node-n2-badge-icon");
+      expect(listenBadge).toBeInTheDocument();
+    });
+
+    it("should not render a badge when it is not present", () => {
+      const nodesWithoutBadges = [testNode("n1", GraphNodeType.Wait, 100, "WaitNode")];
+      render(
+        <div>
+          <RF.ReactFlow nodeTypes={NodeTypes} 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", () => {