diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index 8edc1affa..f6403b8cb 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -14,6 +14,7 @@ import PromptAreaComponent from "../../../../components/promptComponent";
 import TextAreaComponent from "../../../../components/textAreaComponent";
 import ToggleShadComponent from "../../../../components/toggleShadComponent";
 import { MAX_LENGTH_TO_SCROLL_TOOLTIP } from "../../../../constants/constants";
+import { PopUpContext } from "../../../../contexts/popUpContext";
 import { TabsContext } from "../../../../contexts/tabsContext";
 import { typesContext } from "../../../../contexts/typesContext";
 import { ParameterComponentType } from "../../../../types/components";
@@ -27,6 +28,7 @@ import {
   classNames,
   getRandomKeyByssmm,
   groupByFamily,
+  groupByFamilyCustom,
 } from "../../../../utils/utils";
 
 export default function ParameterComponent({
@@ -49,7 +51,8 @@ export default function ParameterComponent({
   const infoHtml = useRef(null);
   const updateNodeInternals = useUpdateNodeInternals();
   const [position, setPosition] = useState(0);
-  const { setTabsState, tabId, save } = useContext(TabsContext);
+  const { setTabsState, tabId, save, flows } = useContext(TabsContext);
+  const { closeEdit } = useContext(PopUpContext);
 
   // Update component position
   useEffect(() => {
@@ -100,56 +103,80 @@ export default function ParameterComponent({
   }, [info]);
 
   useEffect(() => {
-    const groupedObj = groupByFamily(myData, tooltipTitle, left, data.type);
+    let groupedObj = groupByFamily(
+      myData,
+      tooltipTitle,
+      left,
+      data.type,
+      flows.find((f) => f.id === tabId).data.nodes
+    );
 
-    refNumberComponents.current = groupedObj[0]?.type?.length;
+    if (groupedObj?.length === 0) {
+      groupedObj = groupByFamilyCustom(
+        myData,
+        tooltipTitle,
+        left,
+        data.type,
+        flows.find((f) => f.id === tabId).data.nodes
+      );
+    }
 
-    refHtml.current = groupedObj.map((item, i) => {
-      const Icon: any = nodeIconsLucide[item.family];
+    if (groupedObj) {
+      refNumberComponents.current = groupedObj[0]?.type?.length;
 
-      return (
-        <span
-          key={getRandomKeyByssmm() + item.family + i}
-          className={classNames(
-            i > 0 ? "mt-2 flex items-center" : "flex items-center"
-          )}
-        >
-          <div
-            className="h-5 w-5"
-            style={{
-              color: nodeColors[item.family],
-            }}
+      refHtml.current = groupedObj.map((item, i) => {
+        const Icon: any = nodeIconsLucide[item.family];
+
+        return (
+          <span
+            key={getRandomKeyByssmm() + item.family + i}
+            className={classNames(
+              i > 0 ? "mt-2 flex items-center" : "flex items-center"
+            )}
           >
-            <Icon
+            <div
               className="h-5 w-5"
-              strokeWidth={1.5}
               style={{
-                color: nodeColors[item.family] ?? nodeColors.unknown,
+                color: nodeColors[item.family],
               }}
-            />
-          </div>
-          <span className="ps-2 text-xs text-foreground">
-            {nodeNames[item.family] ?? ""}{" "}
-            <span className="text-xs">
-              {" "}
-              {item.type === "" ? "" : " - "}
-              {item.type.split(", ").length > 2
-                ? item.type.split(", ").map((el, i) => (
-                    <React.Fragment key={el + i}>
-                      <span>
-                        {i === item.type.split(", ").length - 1
-                          ? el
-                          : (el += `, `)}
-                      </span>
-                    </React.Fragment>
-                  ))
-                : item.type}
+            >
+              <Icon
+                className="h-5 w-5"
+                strokeWidth={1.5}
+                style={{
+                  color: nodeColors[item.family] ?? nodeColors.unknown,
+                }}
+              />
+            </div>
+            <span className="ps-2 text-xs text-foreground">
+              {item.family !== "custom_components"
+                ? nodeNames[item.family]
+                : item.component ?? ""}{" "}
+              <span className="text-xs">
+                {" "}
+                {item.type === "" ? "" : " - "}
+                {item.type.split(", ").length > 2
+                  ? item.type.split(", ").map((el, i) => (
+                      <React.Fragment key={el + i}>
+                        <span>
+                          {i === item.type.split(", ").length - 1
+                            ? el
+                            : (el += `, `)}
+                        </span>
+                      </React.Fragment>
+                    ))
+                  : item.type}
+              </span>
             </span>
           </span>
-        </span>
-      );
-    });
-  }, [tooltipTitle]);
+        );
+      });
+    }
+  }, [
+    tooltipTitle,
+    flows.find((f) => f.id === tabId).data.nodes.length,
+    closeEdit,
+  ]);
   return (
     <div
       ref={ref}
diff --git a/src/frontend/src/modals/codeAreaModal/index.tsx b/src/frontend/src/modals/codeAreaModal/index.tsx
index 53dda0e81..229396e97 100644
--- a/src/frontend/src/modals/codeAreaModal/index.tsx
+++ b/src/frontend/src/modals/codeAreaModal/index.tsx
@@ -11,9 +11,11 @@ import { Button } from "../../components/ui/button";
 import { CODE_PROMPT_DIALOG_SUBTITLE } from "../../constants/constants";
 import { alertContext } from "../../contexts/alertContext";
 import { darkContext } from "../../contexts/darkContext";
+import { PopUpContext } from "../../contexts/popUpContext";
 import { typesContext } from "../../contexts/typesContext";
 import { postCustomComponent, postValidateCode } from "../../controllers/API";
 import { APIClassType } from "../../types/api";
+import { getRandomKeyByssmm } from "../../utils/utils";
 import BaseModal from "../baseModal";
 
 export default function CodeAreaModal({
@@ -39,6 +41,7 @@ export default function CodeAreaModal({
   const [error, setError] = useState<{
     detail: { error: string; traceback: string };
   }>(null);
+  const { setCloseEdit } = useContext(PopUpContext);
 
   useEffect(() => {
     // if nodeClass.template has more fields other than code and dynamic is true
@@ -61,6 +64,7 @@ export default function CodeAreaModal({
             });
             setOpen(false);
             setValue(code);
+            setCloseEdit(getRandomKeyByssmm().toString());
             // setValue(code);
           } else {
             if (funcErrors.length !== 0) {
@@ -97,6 +101,7 @@ export default function CodeAreaModal({
           setNodeClass(data);
           setValue(code);
           setOpen(false);
+          setCloseEdit(getRandomKeyByssmm().toString());
         }
       })
       .catch((err) => {
diff --git a/src/frontend/src/utils/utils.ts b/src/frontend/src/utils/utils.ts
index 8e6e14fda..d70f6a5a9 100644
--- a/src/frontend/src/utils/utils.ts
+++ b/src/frontend/src/utils/utils.ts
@@ -5,6 +5,7 @@ import { IVarHighlightType } from "../types/components";
 import { FlowType } from "../types/flow";
 import { TabsState } from "../types/tabs";
 import { buildTweaks } from "./reactflowUtils";
+import { nodeNames } from "./styleUtils";
 
 export function classNames(...classes: Array<string>) {
   return classes.filter(Boolean).join(" ");
@@ -88,12 +89,13 @@ export function checkUpperWords(str: string) {
 export const isWrappedWithClass = (event: any, className: string | undefined) =>
   event.target.closest(`.${className}`);
 
-export function groupByFamily(data, baseClasses, left, type) {
+export function groupByFamily(data, baseClasses, left, type, flow) {
   let parentOutput: string;
   let arrOfParent: string[] = [];
   let arrOfType: { family: string; type: string; component: string }[] = [];
   let arrOfLength: { length: number; type: string }[] = [];
   let lastType = "";
+
   Object.keys(data).forEach((d) => {
     Object.keys(data[d]).forEach((n) => {
       try {
@@ -203,6 +205,96 @@ export function groupByFamily(data, baseClasses, left, type) {
   }
 }
 
+export function groupByFamilyCustom(data, baseClasses, left, type, flow) {
+  let arrOfParentCustom: string[] = [];
+  let arrOfType: { family: string; type: string; component: string }[] = [];
+
+  if (type === "CustomComponent") {
+    const uniqueValuesSet = new Set();
+    flow.forEach((element) => {
+      element["data"]["node"]["base_classes"].forEach((el) => {
+        if (!uniqueValuesSet.has(el)) {
+          arrOfParentCustom.push(el);
+          uniqueValuesSet.add(el);
+        }
+      });
+    });
+  }
+
+  if (left === false) {
+    arrOfParentCustom.map((n) => {
+      try {
+        arrOfType.push({
+          family: "custom_components",
+          type: n,
+          component: nodeNames["custom_components"],
+        });
+      } catch (e) {
+        console.log(e);
+      }
+    });
+  } else {
+    flow.forEach((element) => {
+      Object.keys(element["data"]["node"]["template"]).map((el) => {
+        if (
+          element["data"]["node"]["template"][el].input_types &&
+          element["data"]["node"]["template"][el].input_types.length > 0
+        ) {
+          element["data"]["node"]["template"][el].input_types.map((n) => {
+            try {
+              arrOfType.push({
+                family: "custom_components",
+                type: n,
+                component: nodeNames["custom_components"],
+              });
+            } catch (e) {
+              console.log(e);
+            }
+          });
+        }
+      });
+    });
+  }
+
+  const groupedResult = {};
+
+  arrOfType.forEach((item) => {
+    const { family, type, component } = item;
+    if (groupedResult.hasOwnProperty(family)) {
+      if (!groupedResult[family].type.includes(type)) {
+        groupedResult[family].type += `, ${type}`;
+      }
+    } else {
+      groupedResult[family] = { family, type, component };
+    }
+  });
+
+  const result = Object.values(groupedResult);
+
+  if (left === false) {
+    let resultFiltered = [];
+    flow.forEach((element) => {
+      Object.keys(element["data"]["node"]["template"]).map((el) => {
+        if (
+          element["data"]["node"]["template"][el].input_types &&
+          element["data"]["node"]["template"][el].input_types.length > 0
+        ) {
+          element["data"]["node"]["template"][el].input_types.map((n) => {
+            resultFiltered.push({
+              family: "custom_components",
+              type: n,
+              component: element["data"]["node"]["display_name"],
+            });
+          });
+        }
+      });
+    });
+    return resultFiltered;
+  } else {
+    return result;
+  }
+}
+
 export function buildInputs(tabsState, id) {
   return tabsState &&
     tabsState[id] &&