diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index 09e211a1c..8e8657af4 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -37,9 +37,9 @@ import {
   isValidConnection,
   scapedJSONStringfy,
 } from "../../../../utils/reactflowUtils";
-import { nodeColors } from "../../../../utils/styleUtils";
 import {
   classNames,
+  cn,
   groupByFamily,
   isThereModal,
 } from "../../../../utils/utils";
@@ -58,7 +58,7 @@ export default function ParameterComponent({
   data,
   tooltipTitle,
   title,
-  color,
+  colors,
   type,
   name = "",
   required = false,
@@ -282,39 +282,63 @@ export default function ParameterComponent({
             }
             side={left ? "left" : "right"}
           >
-            <Handle
-              data-test-id={`handle-${title.toLowerCase()}-${
-                left ? "target" : "source"
-              }`}
-              type={left ? "target" : "source"}
-              position={left ? Position.Left : Position.Right}
-              key={
-                proxy
-                  ? scapedJSONStringfy({ ...id, proxy })
-                  : scapedJSONStringfy(id)
-              }
-              id={
-                proxy
-                  ? scapedJSONStringfy({ ...id, proxy })
-                  : scapedJSONStringfy(id)
-              }
-              isValidConnection={(connection) =>
-                isValidConnection(connection, nodes, edges)
-              }
-              className={classNames(
-                left ? "my-12 -ml-0.5" : "my-12 -mr-0.5",
-                "h-3 w-3 rounded-full border-2 bg-background",
-                !showNode ? "mt-0" : "",
-              )}
-              style={{
-                borderColor: color ?? nodeColors.unknown,
-              }}
-              onClick={() => {
-                setFilterEdge(
-                  groupByFamily(myData, tooltipTitle!, left, nodes!),
-                );
-              }}
-            ></Handle>
+            <div>
+              <Handle
+                data-test-id={`handle-${title.toLowerCase()}-${
+                  left ? "target" : "source"
+                }`}
+                type={left ? "target" : "source"}
+                position={left ? Position.Left : Position.Right}
+                key={
+                  proxy
+                    ? scapedJSONStringfy({ ...id, proxy })
+                    : scapedJSONStringfy(id)
+                }
+                id={
+                  proxy
+                    ? scapedJSONStringfy({ ...id, proxy })
+                    : scapedJSONStringfy(id)
+                }
+                isValidConnection={(connection) =>
+                  isValidConnection(connection, nodes, edges)
+                }
+                className={classNames(
+                  left ? "my-12 -ml-1" : "my-12 -mr-1",
+                  "h-4 w-4 rounded-full border-2 bg-background",
+                  !showNode ? "mt-0" : "",
+                )}
+                style={{
+                  background:
+                    "conic-gradient(" +
+                    colors
+                      .map(
+                        (color, index) =>
+                          color +
+                          " " +
+                          (360 / colors.length) * index +
+                          "deg " +
+                          (360 / colors.length) * (index + 1) +
+                          "deg",
+                      )
+                      .join(" ,") +
+                    ")",
+                  WebkitMaskImage:
+                    "radial-gradient(transparent 40%, black 44%)",
+                  maskImage: "radial-gradient(transparent 40%, black 44%)",
+                }}
+                onClick={() => {
+                  setFilterEdge(
+                    groupByFamily(myData, tooltipTitle!, left, nodes!),
+                  );
+                }}
+              ></Handle>
+              <div
+                className={cn(
+                  left ? "my-12 -ml-1" : "my-12 -mr-1",
+                  "absolute -left-[4px] top-[50%] z-10 h-4 w-4 translate-y-[-50%] rounded-full bg-background",
+                )}
+              />
+            </div>
           </ShadTooltip>
         </div>
       </Button>
@@ -445,28 +469,54 @@ export default function ParameterComponent({
                 }
                 side={left ? "left" : "right"}
               >
-                <Handle
-                  data-test-id={`handle-${title.toLowerCase()}-${
-                    left ? "left" : "right"
-                  }`}
-                  type={left ? "target" : "source"}
-                  position={left ? Position.Left : Position.Right}
-                  key={scapedJSONStringfy(proxy ? { ...id, proxy } : id)}
-                  id={scapedJSONStringfy(proxy ? { ...id, proxy } : id)}
-                  isValidConnection={(connection) =>
-                    isValidConnection(connection, nodes, edges)
-                  }
-                  className={classNames(
-                    left ? "-ml-0.5" : "-mr-0.5",
-                    "h-3 w-3 rounded-full border-2 bg-background",
-                  )}
-                  style={{ borderColor: color ?? nodeColors.unknown }}
-                  onClick={() => {
-                    setFilterEdge(
-                      groupByFamily(myData, tooltipTitle!, left, nodes!),
-                    );
-                  }}
-                />
+                <div>
+                  <Handle
+                    data-test-id={`handle-${title.toLowerCase()}-${
+                      left ? "left" : "right"
+                    }`}
+                    type={left ? "target" : "source"}
+                    position={left ? Position.Left : Position.Right}
+                    key={scapedJSONStringfy(proxy ? { ...id, proxy } : id)}
+                    id={scapedJSONStringfy(proxy ? { ...id, proxy } : id)}
+                    isValidConnection={(connection) =>
+                      isValidConnection(connection, nodes, edges)
+                    }
+                    className={classNames(
+                      left ? "-ml-1" : "-mr-1",
+                      "z-20 h-4 w-4 rounded-full border-none bg-background",
+                    )}
+                    style={{
+                      background:
+                        "conic-gradient(" +
+                        colors
+                          .map(
+                            (color, index) =>
+                              color +
+                              " " +
+                              (360 / colors.length) * index +
+                              "deg " +
+                              (360 / colors.length) * (index + 1) +
+                              "deg",
+                          )
+                          .join(" ,") +
+                        ")",
+                      WebkitMaskImage:
+                        "radial-gradient(transparent 40%, black 44%)",
+                      maskImage: "radial-gradient(transparent 40%, black 44%)",
+                    }}
+                    onClick={() => {
+                      setFilterEdge(
+                        groupByFamily(myData, tooltipTitle!, left, nodes!),
+                      );
+                    }}
+                  />
+                  <div
+                    className={cn(
+                      "absolute top-[50%] z-10 h-4 w-4 translate-y-[-50%] rounded-full bg-background",
+                      left ? "-left-[4px] -ml-1" : "-right-[4px] -mr-1",
+                    )}
+                  />
+                </div>
               </ShadTooltip>
             </div>
           </Button>
@@ -486,7 +536,7 @@ export default function ParameterComponent({
                   // Commenting this out until we have a better
                   // way to display
                   // (data.node?.template[name]?.refresh ? "w-5/6 " : "") +
-                  "flex-grow mt-2"
+                  "mt-2 flex-grow"
                 }
               >
                 <InputListComponent
diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index ba1b6d77a..3e3da2c65 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -32,6 +32,8 @@ import { nodeColors, nodeIconsLucide } from "../../utils/styleUtils";
 import { classNames, cn } from "../../utils/utils";
 import { countHandlesFn } from "../helpers/count-handles";
 import { getSpecificClassFromBuildStatus } from "../helpers/get-class-from-build-status";
+import { getNodeInputColors } from "../helpers/get-node-input-colors";
+import { getNodeOutputColors } from "../helpers/get-node-output-colors";
 import useCheckCodeValidity from "../hooks/use-check-code-validity";
 import useIconNodeRender from "../hooks/use-icon-render";
 import useIconStatus from "../hooks/use-icons-status";
@@ -279,12 +281,7 @@ export default function GenericNode({
           }) + idx
         }
         data={data}
-        color={
-          nodeColors[output.selected ?? output.types[0]] ??
-          nodeColors[types[output.selected ?? output.types[0]]] ??
-          nodeColors[types[data.type]] ??
-          nodeColors.unknown
-        }
+        colors={getNodeOutputColors(output, data, types)}
         outputProxy={output.proxy}
         title={output.display_name ?? output.name}
         tooltipTitle={output.selected ?? output.types[0]}
@@ -477,36 +474,11 @@ export default function GenericNode({
                               proxy: data.node!.template[templateField].proxy,
                             })}
                             data={data}
-                            color={
-                              data.node?.template[templateField].input_types &&
-                              data.node?.template[templateField].input_types!
-                                .length > 0
-                                ? nodeColors[
-                                    data.node?.template[templateField]
-                                      .input_types![
-                                      data.node?.template[templateField]
-                                        .input_types!.length - 1
-                                    ]
-                                  ] ??
-                                  nodeColors[
-                                    types[
-                                      data.node?.template[templateField]
-                                        .input_types![
-                                        data.node?.template[templateField]
-                                          .input_types!.length - 1
-                                      ]
-                                    ]
-                                  ]
-                                : nodeColors[
-                                    data.node?.template[templateField].type!
-                                  ] ??
-                                  nodeColors[
-                                    types[
-                                      data.node?.template[templateField].type!
-                                    ]
-                                  ] ??
-                                  nodeColors.unknown
-                            }
+                            colors={getNodeInputColors(
+                              data.node?.template[templateField].input_types,
+                              data.node?.template[templateField].type,
+                              types,
+                            )}
                             title={getFieldTitle(
                               data.node?.template!,
                               templateField,
@@ -739,33 +711,11 @@ export default function GenericNode({
                           proxy: data.node!.template[templateField].proxy,
                         })}
                         data={data}
-                        color={
-                          data.node?.template[templateField].input_types &&
-                          data.node?.template[templateField].input_types!
-                            .length > 0
-                            ? nodeColors[
-                                data.node?.template[templateField].input_types![
-                                  data.node?.template[templateField]
-                                    .input_types!.length - 1
-                                ]
-                              ] ??
-                              nodeColors[
-                                types[
-                                  data.node?.template[templateField]
-                                    .input_types![
-                                    data.node?.template[templateField]
-                                      .input_types!.length - 1
-                                  ]
-                                ]
-                              ]
-                            : nodeColors[
-                                data.node?.template[templateField].type!
-                              ] ??
-                              nodeColors[
-                                types[data.node?.template[templateField].type!]
-                              ] ??
-                              nodeColors.unknown
-                        }
+                        colors={getNodeInputColors(
+                          data.node?.template[templateField].input_types,
+                          data.node?.template[templateField].type,
+                          types,
+                        )}
                         title={getFieldTitle(
                           data.node?.template!,
                           templateField,
diff --git a/src/frontend/src/CustomNodes/helpers/get-node-input-colors.ts b/src/frontend/src/CustomNodes/helpers/get-node-input-colors.ts
new file mode 100644
index 000000000..abff92936
--- /dev/null
+++ b/src/frontend/src/CustomNodes/helpers/get-node-input-colors.ts
@@ -0,0 +1,34 @@
+import { nodeColors } from "../../utils/styleUtils";
+
+export function getNodeInputColors(input_types, type, types) {
+  // Helper function to get the color based on type
+  const getColorByType = (type) => nodeColors[type] ?? nodeColors.unknown;
+
+  // If input_types is not null and has elements, map colors based on input_types
+  if (input_types && input_types.length > 0) {
+    // Map through input_types and get colors from nodeColors
+    const colorsFromInputs = input_types
+      .map((input) => nodeColors[input])
+      .filter((color) => color);
+    if (colorsFromInputs.length > 0) {
+      return colorsFromInputs;
+    }
+
+    // If no valid colors found in the previous step, map colors based on types[nodeColors[input]]
+    const colorsFromInputTypes = input_types
+      .map((input) => getColorByType(types[input]))
+      .filter((color) => color);
+    if (colorsFromInputTypes.length > 0) {
+      return colorsFromInputTypes;
+    }
+  }
+
+  // If input_types is null or empty, use the fallback logic
+  const fallbackColors = [getColorByType(type)];
+  if (fallbackColors.length > 0) {
+    return fallbackColors;
+  }
+
+  // Default to unknown color
+  return [nodeColors.unknown];
+}
diff --git a/src/frontend/src/CustomNodes/helpers/get-node-output-colors.ts b/src/frontend/src/CustomNodes/helpers/get-node-output-colors.ts
new file mode 100644
index 000000000..d846a304e
--- /dev/null
+++ b/src/frontend/src/CustomNodes/helpers/get-node-output-colors.ts
@@ -0,0 +1,33 @@
+import { nodeColors } from "../../utils/styleUtils";
+
+export function getNodeOutputColors(output, data, types): string[] {
+  // Helper function to get the color based on type
+  const getColorByType = (type) => nodeColors[type] ?? nodeColors.unknown;
+
+  // Try to get the color based on the selected node
+  let color: string = nodeColors[output.selected];
+  if (color) return [color];
+
+  // Try to get the colors based on the output types
+  let colors: string[] = output.types
+    .map((type) => nodeColors[type])
+    .filter((color) => color);
+  if (colors.length > 0) return colors;
+
+  // Try to get the color based on the type of the selected node
+  color = nodeColors[types[output.selected]];
+  if (color) return [color];
+
+  // Try to get the colors based on the types of output
+  colors = output.types
+    .map((type) => getColorByType(types[type]))
+    .filter((color) => color);
+  if (colors.length > 0) return colors;
+
+  // Try to get the color based on the type in data
+  color = nodeColors[types[data.type]];
+  if (color) return [color];
+
+  // Default to unknown color
+  return [nodeColors.unknown];
+}
diff --git a/src/frontend/src/types/components/index.ts b/src/frontend/src/types/components/index.ts
index b69346fb9..b197b8591 100644
--- a/src/frontend/src/types/components/index.ts
+++ b/src/frontend/src/types/components/index.ts
@@ -65,7 +65,7 @@ export type ParameterComponentType = {
   conditionPath?: string | null;
   key: string;
   id: sourceHandleType | targetHandleType;
-  color: string;
+  colors: string[];
   left: boolean;
   type: string | undefined;
   required?: boolean;