diff --git a/src/frontend/src/contexts/tabsContext.tsx b/src/frontend/src/contexts/tabsContext.tsx
index ba31a5621..7637699e3 100644
--- a/src/frontend/src/contexts/tabsContext.tsx
+++ b/src/frontend/src/contexts/tabsContext.tsx
@@ -16,7 +16,7 @@ import {
 } from "../utils";
 import { alertContext } from "./alertContext";
 import { typesContext } from "./typesContext";
-import { APITemplateType } from "../types/api";
+import { APIClassType, APITemplateType } from "../types/api";
 import ShortUniqueId from "short-unique-id";
 import { addEdge } from "reactflow";
 import {
@@ -192,20 +192,26 @@ export function TabsProvider({ children }: { children: ReactNode }) {
   }
 
   function processFlowEdges(flow) {
+    if(!flow.data || !flow.data.edges) return;
     flow.data.edges.forEach((edge) => {
       edge.className = "";
       edge.style = { stroke: "#555555" };
     });
   }
+  function updateDisplay_name(node:NodeType,template:APIClassType) {
+    node.data.node.display_name = template["display_name"]?template["display_name"]:node.data.type;
+  }
 
   function processFlowNodes(flow) {
-    flow.data.nodes.forEach((node) => {
+    if(!flow.data || !flow.data.nodes) return;
+    flow.data.nodes.forEach((node:NodeType) => {
       const template = templates[node.data.type];
       if (!template) {
         setErrorData({ title: `Unknown node type: ${node.data.type}` });
         return;
       }
       if (Object.keys(template["template"]).length > 0) {
+        updateDisplay_name(node,template);
         updateNodeBaseClasses(node, template);
         updateNodeEdges(flow, node, template);
         updateNodeDescription(node, template);
@@ -214,11 +220,11 @@ export function TabsProvider({ children }: { children: ReactNode }) {
     });
   }
 
-  function updateNodeBaseClasses(node, template) {
+  function updateNodeBaseClasses(node:NodeType,template:APIClassType) {
     node.data.node.base_classes = template["base_classes"];
   }
 
-  function updateNodeEdges(flow, node, template) {
+  function updateNodeEdges(flow:FlowType, node:NodeType,template:APIClassType) {
     flow.data.edges.forEach((edge) => {
       if (edge.source === node.id) {
         edge.sourceHandle = edge.sourceHandle
@@ -230,11 +236,11 @@ export function TabsProvider({ children }: { children: ReactNode }) {
     });
   }
 
-  function updateNodeDescription(node, template) {
+  function updateNodeDescription(node:NodeType,template:APIClassType) {
     node.data.node.description = template["description"];
   }
 
-  function updateNodeTemplate(node, template) {
+  function updateNodeTemplate(node:NodeType,template:APIClassType) {
     node.data.node.template = updateTemplate(
       template["template"] as unknown as APITemplateType,
       node.data.node.template as APITemplateType
@@ -463,6 +469,8 @@ export function TabsProvider({ children }: { children: ReactNode }) {
 
       // Create a new flow with a default name if no flow is provided.
       const newFlow = createNewFlow(flowData, flow);
+      processFlowEdges(newFlow);
+      processFlowNodes(newFlow);
 
       try {
         const { id } = await saveFlowToDatabase(newFlow);
diff --git a/src/frontend/src/types/typesContext/index.ts b/src/frontend/src/types/typesContext/index.ts
index c44bfa4ef..64b6e3e9d 100644
--- a/src/frontend/src/types/typesContext/index.ts
+++ b/src/frontend/src/types/typesContext/index.ts
@@ -1,7 +1,8 @@
 import { ReactFlowInstance } from "reactflow";
+import { APIClassType } from "../api";
 
 const types: { [char: string]: string } = {};
-const template: { [char: string]: string } = {};
+const template: { [char: string]: APIClassType } = {};
 const data: { [char: string]: string } = {};
 
 export type typesContextType = {