From 01885c25de2c6380ec915de358ff6407e1b3e218 Mon Sep 17 00:00:00 2001
From: Ryan Hopper-Lowe <46546486+ryanhopperlowe@users.noreply.github.com>
Date: Thu, 19 Dec 2024 14:58:56 -0600
Subject: [PATCH] chore: remove useForm dependency from BasicToolForm (#987)

* chore: remove useForm dependency from BasicToolForm

useForm returns `useRef().current` under the hood as a hack to optimize rendering, which causes some erratic behavior when certain conditions are met (I think it has to do with resetting the form's default values...????)

Either way, it added way more complexity than it was worth (at least for this component) so removing it was the best solution

* chore: remove redundant add/remove code for Tool Catalog
---
 ui/admin/app/components/agent/ToolForm.tsx    |  32 +++--
 .../app/components/tools/BasicToolForm.tsx    | 123 +++++-------------
 ui/admin/app/components/tools/ToolCatalog.tsx |  24 ++--
 ui/admin/app/components/workflow/Workflow.tsx |   4 +-
 .../app/components/workflow/steps/Step.tsx    |   6 +-
 5 files changed, 71 insertions(+), 118 deletions(-)

diff --git a/ui/admin/app/components/agent/ToolForm.tsx b/ui/admin/app/components/agent/ToolForm.tsx
index 5e7314d7b..f8d1754f4 100644
--- a/ui/admin/app/components/agent/ToolForm.tsx
+++ b/ui/admin/app/components/agent/ToolForm.tsx
@@ -124,6 +124,22 @@ export function ToolForm({
             { tool, variant }
         );
 
+    const updateTools = (tools: string[], variant: ToolVariant) => {
+        const removedToolIndexes = toolFields.fields
+            .filter((field) => !tools.includes(field.tool))
+            .map((item) => toolFields.fields.indexOf(item));
+
+        const addedTools = tools.filter(
+            (tool) => !toolFields.fields.some((field) => field.tool === tool)
+        );
+
+        toolFields.remove(removedToolIndexes);
+
+        for (const tool of addedTools) {
+            toolFields.append({ tool, variant });
+        }
+    };
+
     return (
         <Form {...form}>
             <form
@@ -152,13 +168,9 @@ export function ToolForm({
                 <div className="flex justify-end">
                     <ToolCatalogDialog
                         tools={allTools}
-                        onAddTool={(tool) =>
-                            toolFields.append({
-                                tool,
-                                variant: ToolVariant.FIXED,
-                            })
+                        onUpdateTools={(tools) =>
+                            updateTools(tools, ToolVariant.FIXED)
                         }
-                        onRemoveTools={removeTools}
                     />
                 </div>
 
@@ -212,13 +224,9 @@ export function ToolForm({
                 <div className="flex justify-end">
                     <ToolCatalogDialog
                         tools={allTools}
-                        onAddTool={(tool) =>
-                            toolFields.append({
-                                tool,
-                                variant: ToolVariant.DEFAULT,
-                            })
+                        onUpdateTools={(tools) =>
+                            updateTools(tools, ToolVariant.DEFAULT)
                         }
-                        onRemoveTools={removeTools}
                         className="w-auto"
                     />
                 </div>
diff --git a/ui/admin/app/components/tools/BasicToolForm.tsx b/ui/admin/app/components/tools/BasicToolForm.tsx
index ae7bc2588..512f2aaa7 100644
--- a/ui/admin/app/components/tools/BasicToolForm.tsx
+++ b/ui/admin/app/components/tools/BasicToolForm.tsx
@@ -1,105 +1,48 @@
-import { zodResolver } from "@hookform/resolvers/zod";
-import { useEffect, useMemo } from "react";
-import { useFieldArray, useForm } from "react-hook-form";
-import { z } from "zod";
+import { memo, useCallback, useMemo, useState } from "react";
 
 import { ToolEntry } from "~/components/agent/ToolEntry";
 import { ToolCatalogDialog } from "~/components/tools/ToolCatalog";
-import { Form } from "~/components/ui/form";
 
-const formSchema = z.object({
-    tools: z.array(z.object({ value: z.string() })),
-});
-
-type BasicToolFormValues = z.infer<typeof formSchema>;
-
-type Tools = { tools: string[] };
-
-export function BasicToolForm({
-    defaultValues: _defaultValues,
-    onChange,
-}: {
-    defaultValues?: Partial<Tools>;
-    onSubmit?: (values: Tools) => void;
-    onChange?: (values: Tools) => void;
+export const BasicToolForm = memo(function BasicToolFormComponent(props: {
+    value?: string[];
+    defaultValue?: string[];
+    onChange?: (values: string[]) => void;
 }) {
-    const defaultValues = useMemo(() => {
-        return {
-            tools:
-                _defaultValues?.tools?.map((tool) => ({ value: tool })) || [],
-        };
-    }, [_defaultValues]);
-
-    const form = useForm<BasicToolFormValues>({
-        resolver: zodResolver(formSchema),
-        defaultValues: { tools: defaultValues?.tools || [] },
-    });
-    const { getValues, reset } = form;
-
-    useEffect(() => {
-        const unchanged = compareArrays(
-            defaultValues?.tools.map(({ value }) => value) || [],
-            getValues().tools.map(({ value }) => value)
-        );
+    const { onChange } = props;
 
-        if (unchanged) return;
-
-        reset({ tools: defaultValues?.tools || [] });
-    }, [defaultValues, getValues, reset]);
-
-    const toolArr = useFieldArray({ control: form.control, name: "tools" });
-
-    useEffect(() => {
-        return form.watch((values) => {
-            const { data, success } = formSchema.safeParse(values);
-
-            if (!success) return;
+    const [_value, _setValue] = useState(props.defaultValue);
+    const value = useMemo(
+        () => props.value ?? _value ?? [],
+        [props.value, _value]
+    );
 
-            onChange?.({ tools: data.tools.map((t) => t.value) });
-        }).unsubscribe;
-    }, [form, onChange]);
+    const setValue = useCallback(
+        (newValue: string[]) => {
+            _setValue(newValue);
+            onChange?.(newValue);
+        },
+        [onChange]
+    );
 
     const removeTools = (toolsToRemove: string[]) => {
-        const indexes = toolsToRemove
-            .map((tool) => toolArr.fields.findIndex((t) => t.value === tool))
-            .filter((index) => index !== -1);
-
-        toolArr.remove(indexes);
-    };
-
-    const addTool = (tool: string) => {
-        toolArr.append({ value: tool });
+        setValue(value.filter((tool) => !toolsToRemove.includes(tool)));
     };
 
     return (
-        <Form {...form}>
-            <div className="flex flex-col gap-2">
-                <div className="flex flex-col gap-1 w-full overflow-y-auto">
-                    {toolArr.fields.map((field) => (
-                        <ToolEntry
-                            key={field.id}
-                            tool={field.value}
-                            onDelete={() => removeTools([field.value])}
-                        />
-                    ))}
-                </div>
-
-                <div className="flex justify-end">
-                    <ToolCatalogDialog
-                        tools={toolArr.fields.map((field) => field.value)}
-                        onAddTool={addTool}
-                        onRemoveTools={removeTools}
+        <div className="flex flex-col gap-2">
+            <div className="flex flex-col gap-1 w-full overflow-y-auto">
+                {value.map((tool) => (
+                    <ToolEntry
+                        key={tool}
+                        tool={tool}
+                        onDelete={() => removeTools([tool])}
                     />
-                </div>
+                ))}
             </div>
-        </Form>
-    );
-}
 
-function compareArrays(a: string[], b: string[]) {
-    const aSet = new Set(a);
-
-    if (aSet.size !== b.length) return false;
-
-    return b.every((tool) => aSet.has(tool));
-}
+            <div className="flex justify-end">
+                <ToolCatalogDialog tools={value} onUpdateTools={setValue} />
+            </div>
+        </div>
+    );
+});
diff --git a/ui/admin/app/components/tools/ToolCatalog.tsx b/ui/admin/app/components/tools/ToolCatalog.tsx
index b8441f713..60eeecab2 100644
--- a/ui/admin/app/components/tools/ToolCatalog.tsx
+++ b/ui/admin/app/components/tools/ToolCatalog.tsx
@@ -27,8 +27,7 @@ import {
 
 type ToolCatalogProps = React.HTMLAttributes<HTMLDivElement> & {
     tools: string[];
-    onAddTool: (tools: string) => void;
-    onRemoveTools: (tools: string[]) => void;
+    onUpdateTools: (tools: string[]) => void;
     invert?: boolean;
     classNames?: { list?: string };
 };
@@ -37,8 +36,7 @@ export function ToolCatalog({
     className,
     tools,
     invert = false,
-    onAddTool,
-    onRemoveTools,
+    onUpdateTools,
     classNames,
 }: ToolCatalogProps) {
     const { data: toolCategories, isLoading } = useSWR(
@@ -50,25 +48,29 @@ export function ToolCatalog({
     const handleSelect = useCallback(
         (toolId: string) => {
             if (!tools.includes(toolId)) {
-                onAddTool(toolId);
+                onUpdateTools([...tools, toolId]);
             }
         },
-        [tools, onAddTool]
+        [tools, onUpdateTools]
     );
 
     const handleSelectBundle = useCallback(
         (bundleToolId: string, categoryTools: ToolReference[]) => {
             if (tools.includes(bundleToolId)) {
-                onRemoveTools([bundleToolId]);
+                onUpdateTools(tools.filter((tool) => tool !== bundleToolId));
                 return;
             }
 
-            onAddTool(bundleToolId);
+            const toolsToRemove = new Set(categoryTools.map((tool) => tool.id));
 
-            // remove all tools in the bundle to remove redundancy
-            onRemoveTools(categoryTools.map((tool) => tool.id));
+            const newTools = [
+                ...tools.filter((tool) => !toolsToRemove.has(tool)),
+                bundleToolId,
+            ];
+
+            onUpdateTools(newTools);
         },
-        [tools, onAddTool, onRemoveTools]
+        [tools, onUpdateTools]
     );
 
     if (isLoading) return <LoadingSpinner />;
diff --git a/ui/admin/app/components/workflow/Workflow.tsx b/ui/admin/app/components/workflow/Workflow.tsx
index 058dfccfa..62546c4ed 100644
--- a/ui/admin/app/components/workflow/Workflow.tsx
+++ b/ui/admin/app/components/workflow/Workflow.tsx
@@ -78,8 +78,8 @@ function WorkflowContent({ className }: WorkflowProps) {
                     </CardDescription>
 
                     <BasicToolForm
-                        defaultValues={workflow}
-                        onChange={debouncedSetWorkflowInfo}
+                        value={workflow.tools}
+                        onChange={(tools) => partialSetWorkflow({ tools })}
                     />
                 </div>
 
diff --git a/ui/admin/app/components/workflow/steps/Step.tsx b/ui/admin/app/components/workflow/steps/Step.tsx
index 38d465e8c..80d487aa8 100644
--- a/ui/admin/app/components/workflow/steps/Step.tsx
+++ b/ui/admin/app/components/workflow/steps/Step.tsx
@@ -107,9 +107,9 @@ export function StepComponent({
 
                             <AccordionContent className="p-1 pb-6">
                                 <BasicToolForm
-                                    defaultValues={step}
-                                    onChange={(values) =>
-                                        onUpdate({ ...step, ...values })
+                                    value={step.tools}
+                                    onChange={(tools) =>
+                                        onUpdate({ ...step, tools })
                                     }
                                 />
                             </AccordionContent>