From 4ca4c26f5a759a78b4dd24b47c7ec89fb1f430a9 Mon Sep 17 00:00:00 2001
From: Oliver Walters <oliver.henry.walters@gmail.com>
Date: Sun, 21 Jun 2026 03:47:51 +0000
Subject: [PATCH] Custom node rendering

---
 .../src/components/forms/fields/TreeField.tsx | 81 ++++++++++++++++++-
 1 file changed, 80 insertions(+), 1 deletion(-)

diff --git a/src/frontend/src/components/forms/fields/TreeField.tsx b/src/frontend/src/components/forms/fields/TreeField.tsx
index d9255c5766..64f7bc5fd8 100644
--- a/src/frontend/src/components/forms/fields/TreeField.tsx
+++ b/src/frontend/src/components/forms/fields/TreeField.tsx
@@ -1,12 +1,14 @@
 import { t } from '@lingui/core/macro';
-import { type TreeNodeData, TreeSelect } from '@mantine/core';
+import { Group, Text, type TreeNodeData, TreeSelect } from '@mantine/core';
 import { useDebouncedValue } from '@mantine/hooks';
+import { IconChevronDown, IconChevronRight } from '@tabler/icons-react';
 import { useQuery } from '@tanstack/react-query';
 import { useCallback, useEffect, useMemo, useState } from 'react';
 import type { FieldValues, UseControllerReturn } from 'react-hook-form';
 
 import type { ApiEndpoints } from '@lib/enums/ApiEndpoints';
 import { apiUrl } from '@lib/functions/Api';
+import { cancelEvent } from '@lib/functions/Events';
 import type { ApiFormFieldType } from '@lib/types/Forms';
 import { useApi } from '../../../contexts/ApiContext';
 
@@ -51,6 +53,13 @@ export function TreeField({
 
   const nodes: any[] = useMemo(() => query.data ?? [], [query.data]);
 
+  // O(1) lookup for raw node data inside renderNode
+  const nodeMap = useMemo(() => {
+    const map: Record<string, any> = {};
+    for (const n of nodes) map[n.pk.toString()] = n;
+    return map;
+  }, [nodes]);
+
   // Expand all returned nodes when a search is active so users can see all matches.
   // Collapse back to root when the search is cleared.
   useEffect(() => {
@@ -124,6 +133,76 @@ export function TreeField({
       nothingFoundMessage={
         query.isFetching ? t`Loading...` : t`No results found`
       }
+      renderNode={({ node, expanded, hasChildren, selected }) => {
+        const raw = nodeMap[node.value];
+        return (
+          <Group
+            justify='space-between'
+            gap='xs'
+            wrap='nowrap'
+            style={{ flex: 1 }}
+          >
+            <Group gap={4} wrap='nowrap'>
+              {/* Chevron rendered manually so renderNode can coexist with expand behavior.
+                  stopPropagation prevents the Combobox.Option from selecting the node
+                  when the user clicks the expand toggle. */}
+              <span
+                style={{
+                  display: 'inline-flex',
+                  width: 16,
+                  alignItems: 'center',
+                  flexShrink: 0,
+                  cursor: hasChildren ? 'pointer' : 'default'
+                }}
+                role={hasChildren ? 'button' : undefined}
+                tabIndex={hasChildren ? -1 : undefined}
+                aria-label={expanded ? t`Collapse` : t`Expand`}
+                onClick={
+                  hasChildren
+                    ? (event: any) => {
+                        cancelEvent(event);
+                        setExpandedValues((prev) =>
+                          prev.includes(node.value)
+                            ? prev.filter((v) => v !== node.value)
+                            : [...prev, node.value]
+                        );
+                      }
+                    : undefined
+                }
+                onKeyDown={
+                  hasChildren
+                    ? (event: any) => {
+                        if (event.key === 'Enter' || event.key === ' ') {
+                          cancelEvent(event);
+                          setExpandedValues((prev) =>
+                            prev.includes(node.value)
+                              ? prev.filter((v) => v !== node.value)
+                              : [...prev, node.value]
+                          );
+                        }
+                      }
+                    : undefined
+                }
+              >
+                {hasChildren &&
+                  (expanded ? (
+                    <IconChevronDown size={14} />
+                  ) : (
+                    <IconChevronRight size={14} />
+                  ))}
+              </span>
+              <Text size='sm' fw={selected ? 600 : undefined}>
+                {raw?.name ?? String(node.label)}
+              </Text>
+            </Group>
+            {raw?.description && (
+              <Text size='xs' c='dimmed' ta='right' truncate flex={1} maw='50%'>
+                {raw.description}
+              </Text>
+            )}
+          </Group>
+        );
+      }}
     />
   );
 }