From 309c5cc6aa49594c103d42f2181101b51dc71ca4 Mon Sep 17 00:00:00 2001
From: tylerslaton <mtslaton1@gmail.com>
Date: Thu, 24 Oct 2024 10:11:11 -0400
Subject: [PATCH] fix: use same getter as ToolCatalog in ToolGrid

Signed-off-by: tylerslaton <mtslaton1@gmail.com>
---
 .../components/tools/toolGrid/ToolGrid.tsx    | 80 ++++++++++---------
 .../lib/service/api/toolreferenceService.ts   |  9 +--
 ui/admin/app/routes/_auth.tools._index.tsx    | 15 ++--
 3 files changed, 51 insertions(+), 53 deletions(-)

diff --git a/ui/admin/app/components/tools/toolGrid/ToolGrid.tsx b/ui/admin/app/components/tools/toolGrid/ToolGrid.tsx
index f9f08f4df..fe73b5128 100644
--- a/ui/admin/app/components/tools/toolGrid/ToolGrid.tsx
+++ b/ui/admin/app/components/tools/toolGrid/ToolGrid.tsx
@@ -1,57 +1,59 @@
 import { useMemo } from "react";
 
 import { ToolReference } from "~/lib/model/toolReferences";
+import { ToolCategoryMap } from "~/lib/service/api/toolreferenceService";
 
 import { CategoryHeader } from "~/components/tools/toolGrid/CategoryHeader";
 import { CategoryTools } from "~/components/tools/toolGrid/CategoryTools";
 
 interface ToolGridProps {
-    tools: ToolReference[];
+    toolCategories: ToolCategoryMap;
     filter: string;
     onDelete: (id: string) => void;
 }
 
-export function ToolGrid({ tools, filter, onDelete }: ToolGridProps) {
-    const filteredTools = useMemo(() => {
-        return tools?.filter(
-            (tool) =>
-                tool.name?.toLowerCase().includes(filter.toLowerCase()) ||
-                tool.metadata?.category
-                    ?.toLowerCase()
-                    .includes(filter.toLowerCase()) ||
-                tool.description?.toLowerCase().includes(filter.toLowerCase())
-        );
-    }, [tools, filter]);
-
-    const sortedCategories = useMemo(() => {
-        const categorizedTools = filteredTools?.reduce(
-            (acc, tool) => {
-                const category = tool.metadata?.category ?? "Uncategorized";
-                if (!acc[category]) {
-                    acc[category] = [];
-                }
-                acc[category].push(tool);
-                return acc;
-            },
-            {} as Record<string, ToolReference[]>
-        );
-
-        // Sort categories to put "Uncategorized" first
-        return Object.entries(categorizedTools).sort(([a], [b]) => {
-            if (a === "Uncategorized") return -1;
-            if (b === "Uncategorized") return 1;
-            return a.localeCompare(b);
-        });
-    }, [filteredTools]);
+export function ToolGrid({ toolCategories, filter, onDelete }: ToolGridProps) {
+    const filteredCategories = useMemo(() => {
+        const result: ToolCategoryMap = {};
+        for (const [category, { tools, bundleTool }] of Object.entries(
+            toolCategories
+        )) {
+            const filteredTools = tools.filter(
+                (tool) =>
+                    tool.name?.toLowerCase().includes(filter.toLowerCase()) ||
+                    tool.metadata?.category
+                        ?.toLowerCase()
+                        .includes(filter.toLowerCase()) ||
+                    tool.description
+                        ?.toLowerCase()
+                        .includes(filter.toLowerCase())
+            );
+            if (filteredTools.length > 0 || bundleTool) {
+                result[category] = {
+                    tools: filteredTools,
+                    bundleTool: bundleTool,
+                };
+            }
+        }
+        return result;
+    }, [toolCategories, filter]);
 
     return (
         <div className="space-y-8 pb-16">
-            {sortedCategories.map(([category, categoryTools]) => (
-                <div key={category} className="space-y-4">
-                    <CategoryHeader category={category} tools={categoryTools} />
-                    <CategoryTools tools={categoryTools} onDelete={onDelete} />
-                </div>
-            ))}
+            {Object.entries(filteredCategories).map(
+                ([category, { tools, bundleTool }]) => (
+                    <div key={category} className="space-y-4">
+                        <CategoryHeader category={category} tools={tools} />
+                        <CategoryTools
+                            tools={[
+                                bundleTool || ({} as ToolReference),
+                                ...tools,
+                            ]}
+                            onDelete={onDelete}
+                        />
+                    </div>
+                )
+            )}
         </div>
     );
 }
diff --git a/ui/admin/app/lib/service/api/toolreferenceService.ts b/ui/admin/app/lib/service/api/toolreferenceService.ts
index 6e3d4bb44..103c6cfb2 100644
--- a/ui/admin/app/lib/service/api/toolreferenceService.ts
+++ b/ui/admin/app/lib/service/api/toolreferenceService.ts
@@ -54,13 +54,8 @@ async function getToolReferencesCategoryMap(type?: ToolReferenceType) {
 }
 getToolReferencesCategoryMap.key = (type?: ToolReferenceType) =>
     ({
-        // tylerslaton: This is a workaround to make the SWR caching work
-        // nicely with the getToolReferences function. The reason we do this
-        // is because otherwise the SWR cache will not be invalidated between
-        // the two function calls and thus break type safety since types in
-        // TypeScript are erased at runtime. as-map=true does not exist in
-        // the API.
-        url: ApiRoutes.toolReferences.base({ type }).path + "?as-map=true",
+        url: ApiRoutes.toolReferences.base({ type }).path,
+        responseType: "map",
     }) as const;
 
 const getToolReferenceById = async (toolReferenceId: string) => {
diff --git a/ui/admin/app/routes/_auth.tools._index.tsx b/ui/admin/app/routes/_auth.tools._index.tsx
index e4866b869..2cae3f387 100644
--- a/ui/admin/app/routes/_auth.tools._index.tsx
+++ b/ui/admin/app/routes/_auth.tools._index.tsx
@@ -19,17 +19,18 @@ import { Input } from "~/components/ui/input";
 
 export async function clientLoader() {
     await Promise.all([
-        preload(ToolReferenceService.getToolReferences.key("tool"), () =>
-            ToolReferenceService.getToolReferences("tool")
+        preload(ToolReferenceService.getToolReferencesCategoryMap.key("tool"), () =>
+            ToolReferenceService.getToolReferencesCategoryMap("tool")
         ),
     ]);
     return null;
 }
 
 export default function Tools() {
-    const { data: tools, mutate } = useSWR(
-        ToolReferenceService.getToolReferences.key("tool"),
-        () => ToolReferenceService.getToolReferences("tool")
+    const { data: toolCategories, mutate } = useSWR(
+        ToolReferenceService.getToolReferencesCategoryMap.key("tool"),
+        () => ToolReferenceService.getToolReferencesCategoryMap("tool"),
+        { fallbackData: {} }
     );
 
     const [isDialogOpen, setIsDialogOpen] = useState(false);
@@ -79,9 +80,9 @@ export default function Tools() {
                 </div>
             </div>
 
-            {tools && (
+            {toolCategories && (
                 <ToolGrid
-                    tools={tools}
+                    toolCategories={toolCategories}
                     filter={searchQuery}
                     onDelete={handleDelete}
                 />