diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index 9a1aaf7f2..18c2bbb6a 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -232,9 +232,9 @@ export default function GenericNode({
     if (data.node) {
       postCustomComponent(currentCode, data.node)
         .then((apiReturn) => {
-          const { data } = apiReturn;
-          if (data && updateNodeCode) {
-            updateNodeCode(data, currentCode, "code");
+          const { data, type } = apiReturn.data;
+          if (data && type && updateNodeCode) {
+            updateNodeCode(data, currentCode, "code", type);
             setLoadingUpdate(false);
           }
         })
diff --git a/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx b/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx
index 6fb78ce10..933f836a4 100644
--- a/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx
+++ b/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx
@@ -7,7 +7,7 @@ const useHandleNodeClass = (
   setNode,
   updateNodeInternals,
 ) => {
-  const handleNodeClass = (newNodeClass, code) => {
+  const handleNodeClass = (newNodeClass, code, type?: string) => {
     if (!data.node) return;
     if (data.node!.template[name].value !== code) {
       takeSnapshot();
@@ -22,7 +22,9 @@ const useHandleNodeClass = (
         description: newNodeClass.description ?? data.node!.description,
         display_name: newNodeClass.display_name ?? data.node!.display_name,
       };
-
+      if (type) {
+        newNode.data.node.template[name].type = type;
+      }
       newNode.data.node.template[name].value = code;
 
       return newNode;
diff --git a/src/frontend/src/CustomNodes/hooks/use-update-node-code.tsx b/src/frontend/src/CustomNodes/hooks/use-update-node-code.tsx
index d6c972ddf..1e8e79a1f 100644
--- a/src/frontend/src/CustomNodes/hooks/use-update-node-code.tsx
+++ b/src/frontend/src/CustomNodes/hooks/use-update-node-code.tsx
@@ -11,7 +11,7 @@ const useUpdateNodeCode = (
   updateNodeInternals: (id: string) => void,
 ) => {
   const updateNodeCode = useCallback(
-    (newNodeClass: APIClassType, code: string, name: string) => {
+    (newNodeClass: APIClassType, code: string, name: string, type: string) => {
       setNode(dataId, (oldNode) => {
         let newNode = cloneDeep(oldNode);
 
@@ -22,6 +22,9 @@ const useUpdateNodeCode = (
           display_name: newNodeClass.display_name ?? dataNode.display_name,
           edited: false,
         };
+        if (type) {
+          newNode.data.type = type;
+        }
 
         newNode.data.node.template[name].value = code;
         setIsOutdated(false);
diff --git a/src/frontend/src/controllers/API/index.ts b/src/frontend/src/controllers/API/index.ts
index b20a4c23b..8a4bff599 100644
--- a/src/frontend/src/controllers/API/index.ts
+++ b/src/frontend/src/controllers/API/index.ts
@@ -7,6 +7,7 @@ import {
   APIObjectType,
   APITemplateType,
   Component,
+  CustomComponentRequest,
   LoginType,
   ProfilePicturesTypeAPI,
   Users,
@@ -386,7 +387,7 @@ export async function getProfilePictures(): Promise<ProfilePicturesTypeAPI | nul
 export async function postCustomComponent(
   code: string,
   apiClass: APIClassType,
-): Promise<AxiosResponse<APIClassType>> {
+): Promise<AxiosResponse<CustomComponentRequest>> {
   // let template = apiClass.template;
   return await api.post(`${BASE_URL_API}custom_component`, {
     code,
diff --git a/src/frontend/src/modals/codeAreaModal/index.tsx b/src/frontend/src/modals/codeAreaModal/index.tsx
index 7b3f04c09..74ced8078 100644
--- a/src/frontend/src/modals/codeAreaModal/index.tsx
+++ b/src/frontend/src/modals/codeAreaModal/index.tsx
@@ -102,9 +102,9 @@ export default function CodeAreaModal({
   function processDynamicField() {
     postCustomComponent(code, nodeClass!)
       .then((apiReturn) => {
-        const { data } = apiReturn;
-        if (data) {
-          setNodeClass(data, code);
+        const { data, type } = apiReturn.data;
+        if (data && type) {
+          setNodeClass(data, code, type);
           setError({ detail: { error: undefined, traceback: undefined } });
           setOpen(false);
         }
diff --git a/src/frontend/src/types/api/index.ts b/src/frontend/src/types/api/index.ts
index 50b76737a..f5f354805 100644
--- a/src/frontend/src/types/api/index.ts
+++ b/src/frontend/src/types/api/index.ts
@@ -13,6 +13,11 @@ export type CustomFieldsType = {
   [key: string]: Array<string>;
 };
 
+export type CustomComponentRequest = {
+  data: APIClassType;
+  type: string;
+};
+
 export type APIClassType = {
   base_classes?: Array<string>;
   description: string;
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index 31860309d..992c012bd 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -595,7 +595,11 @@ export type codeAreaModalPropsType = {
   setOpenModal?: (bool: boolean) => void;
   value: string;
   nodeClass: APIClassType | undefined;
-  setNodeClass: (Class: APIClassType, code?: string) => void | undefined;
+  setNodeClass: (
+    Class: APIClassType,
+    code?: string,
+    type?: string,
+  ) => void | undefined;
   children: ReactNode;
   dynamic?: boolean;
   readonly?: boolean;