added suspense to index page (tensorzero#3102)

virajmehta · web-flow · commit f338fc1cfaa8 · 2025-08-20T16:28:41.000Z
* added suspense to index page

* fixed improper memoization of gateway connection info
diff --git a/ui/app/components/layout/TensorZeroStatusIndicator.tsx b/ui/app/components/layout/TensorZeroStatusIndicator.tsx
@@ -1,5 +1,6 @@
 import { useTensorZeroStatusFetcher } from "~/routes/api/tensorzero/status";
 import { Tooltip, TooltipContent, TooltipTrigger } from "../ui/tooltip";
+import { useMemo } from "react";
 
 /**
  * A component that displays the status of the TensorZero Gateway.
@@ -11,26 +12,26 @@ export default function TensorZeroStatusIndicator() {
   // Extract server version from status if available
   const serverVersion = status?.version || "";
 
-  // Check if versions match (ignoring patch version)
-  const versionsMatch = () => {
+  // Check if versions match (ignoring patch version) - memoized to prevent re-renders
+  const versionsMatch = useMemo(() => {
     if (!serverVersion) return true; // No data yet, don't show warning
     // We can do an exact match for now
     return uiVersion === serverVersion;
-  };
+  }, [uiVersion, serverVersion]);
 
-  const getStatusColor = () => {
+  const statusColor = useMemo(() => {
     if (isLoading || status === undefined) return "bg-gray-300"; // Loading or initial state
     if (!status) return "bg-red-500"; // Could not connect (explicit null/failed state)
-    if (!versionsMatch()) return "bg-yellow-500"; // Version mismatch
+    if (!versionsMatch) return "bg-yellow-500"; // Version mismatch
     return "bg-green-500"; // Everything is good
-  };
+  }, [isLoading, status, versionsMatch]);
 
   return (
     <div className="px-3 py-2 text-xs">
       <div className="text-fg-muted flex flex-col gap-1 truncate">
         <div className="flex items-center gap-2">
           <div
-            className={`h-2 w-2 rounded-full ${getStatusColor()} mr-1 inline-block`}
+            className={`h-2 w-2 rounded-full ${statusColor} mr-1 inline-block`}
           />
           {isLoading
             ? "Checking status..."
@@ -40,7 +41,7 @@ export default function TensorZeroStatusIndicator() {
                 ? `TensorZero Gateway ${serverVersion}`
                 : "Failed to connect to Gateway"}
         </div>
-        {status && !versionsMatch() && (
+        {status && !versionsMatch && (
           <div className="ml-5">
             <Tooltip>
               <TooltipTrigger asChild>
diff --git a/ui/app/components/layout/app.sidebar.tsx b/ui/app/components/layout/app.sidebar.tsx
@@ -104,7 +104,7 @@ const navigation: NavigationSection[] = [
 
 export function AppSidebar({ ...props }: React.ComponentProps<typeof Sidebar>) {
   const { state } = useSidebar();
-  const isActivePath = useActivePath();
+  const activePathUtils = useActivePath();
 
   return (
     <Sidebar collapsible="icon" {...props}>
@@ -132,7 +132,7 @@ export function AppSidebar({ ...props }: React.ComponentProps<typeof Sidebar>) {
               <SidebarMenuButton
                 asChild
                 tooltip={state === "collapsed" ? "Dashboard" : undefined}
-                isActive={isActivePath("/")}
+                isActive={activePathUtils.isActive("/")}
               >
                 <Link to="/" className="flex items-center gap-2">
                   <Home className="h-4 w-4" />
@@ -153,7 +153,7 @@ export function AppSidebar({ ...props }: React.ComponentProps<typeof Sidebar>) {
                   <SidebarMenuButton
                     asChild
                     tooltip={state === "collapsed" ? item.title : undefined}
-                    isActive={isActivePath(item.url)}
+                    isActive={activePathUtils.isActive(item.url)}
                   >
                     <Link
                       to={item.url}
diff --git a/ui/app/hooks/use-active-path.ts b/ui/app/hooks/use-active-path.ts
@@ -1,12 +1,21 @@
 import { useLocation } from "react-router";
+import { useMemo } from "react";
 
 export function useActivePath() {
   const location = useLocation();
+  const pathname = location.pathname;
 
-  return (path: string) => {
-    if (path === "/") {
-      return location.pathname === "/";
-    }
-    return location.pathname.startsWith(path);
-  };
+  // Return a stable object with the function as a property
+  // This avoids creating new function references
+  return useMemo(
+    () => ({
+      isActive: (path: string) => {
+        if (path === "/") {
+          return pathname === "/";
+        }
+        return pathname.startsWith(path);
+      },
+    }),
+    [pathname],
+  );
 }
diff --git a/ui/app/routes/api/tensorzero/status.ts b/ui/app/routes/api/tensorzero/status.ts
@@ -1,6 +1,6 @@
 import { useFetcher } from "react-router";
 import { getTensorZeroClient } from "~/utils/tensorzero.server";
-import { useEffect } from "react";
+import { useEffect, useMemo } from "react";
 import { logger } from "~/utils/logger";
 
 export async function loader() {
@@ -20,14 +20,14 @@ export async function loader() {
 export function useTensorZeroStatusFetcher() {
   const statusFetcher = useFetcher();
   const status = statusFetcher.data;
+  const isLoading = statusFetcher.state === "loading";
 
   useEffect(() => {
     if (statusFetcher.state === "idle" && !statusFetcher.data) {
       statusFetcher.load("/api/tensorzero/status");
     }
-    // TODO: Fix and stop ignoring lint rule
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, []);
+  }, []); // Empty dependency array - only run on mount
 
-  return { status, isLoading: statusFetcher.state === "loading" };
+  return useMemo(() => ({ status, isLoading }), [status, isLoading]);
 }
diff --git a/ui/app/routes/index.tsx b/ui/app/routes/index.tsx
@@ -1,4 +1,5 @@
-import { Link, type RouteHandle } from "react-router";
+import { Link, type RouteHandle, Await } from "react-router";
+import * as React from "react";
 import { Card } from "~/components/ui/card";
 import { PageLayout } from "~/components/layout/PageLayout";
 import {
@@ -21,10 +22,9 @@ import {
   countInferencesByFunction,
   countEpisodes,
 } from "~/utils/clickhouse/inference.server";
-import { getAllFunctionConfigs } from "~/utils/config/index.server";
+import { getConfig, getAllFunctionConfigs } from "~/utils/config/index.server";
 import { getDatasetCounts } from "~/utils/clickhouse/datasets.server";
 import { countTotalEvaluationRuns } from "~/utils/clickhouse/evaluations.server";
-import { useConfig } from "~/context/config";
 import type { Route } from "./+types/index";
 import {
   countDynamicEvaluationProjects,
@@ -39,7 +39,7 @@ interface DirectoryCardProps {
   source: string;
   icon: React.ComponentType<{ className?: string; size?: number }>;
   title: string;
-  description: string;
+  description: string | Promise<string>;
 }
 
 function DirectoryCard({
@@ -62,7 +62,19 @@ function DirectoryCard({
             {title}
           </h3>
           <p className="text-fg-secondary overflow-hidden text-xs text-ellipsis whitespace-nowrap">
-            {description}
+            {typeof description === "string" ? (
+              description
+            ) : (
+              <React.Suspense
+                fallback={
+                  <span className="bg-bg-tertiary inline-block h-3 w-16 animate-pulse rounded"></span>
+                }
+              >
+                <Await resolve={description}>
+                  {(resolvedDescription) => resolvedDescription}
+                </Await>
+              </React.Suspense>
+            )}
           </p>
         </div>
       </Card>
@@ -93,55 +105,88 @@ function FooterLink({ source, icon: Icon, children }: FooterLinkProps) {
 }
 
 export async function loader() {
-  const [
-    countsInfo,
-    numEpisodes,
-    datasetCounts,
-    numEvaluationRuns,
-    numDynamicEvaluationRuns,
-    numDynamicEvaluationRunProjects,
-    functionConfigs,
-  ] = await Promise.all([
-    countInferencesByFunction(),
-    countEpisodes(),
-    getDatasetCounts({}),
-    countTotalEvaluationRuns(),
-    countDynamicEvaluationRuns(),
-    countDynamicEvaluationProjects(),
-    getAllFunctionConfigs(),
-  ]);
-  const totalInferences = countsInfo.reduce((acc, curr) => acc + curr.count, 0);
-  const numFunctions = Object.keys(functionConfigs).length;
-  const numVariants = Object.values(functionConfigs).reduce((acc, config) => {
-    return acc + (config ? Object.keys(config.variants || {}).length : 0);
-  }, 0);
-  const numDatasets = datasetCounts.length;
+  // Create the promises
+  const countsInfoPromise = countInferencesByFunction();
+  const numEpisodesPromise = countEpisodes();
+  const datasetCountsPromise = getDatasetCounts({});
+  const numEvaluationRunsPromise = countTotalEvaluationRuns();
+  const numDynamicEvaluationRunsPromise = countDynamicEvaluationRuns();
+  const numDynamicEvaluationRunProjectsPromise =
+    countDynamicEvaluationProjects();
+  const configPromise = getConfig();
+  const functionConfigsPromise = getAllFunctionConfigs();
+
+  // Create derived promises - these will be stable references
+  const totalInferencesDesc = countsInfoPromise.then((countsInfo) => {
+    const total = countsInfo.reduce((acc, curr) => acc + curr.count, 0);
+    return `${total.toLocaleString()} inferences`;
+  });
+
+  const numFunctionsDesc = functionConfigsPromise.then((functionConfigs) => {
+    const numFunctions = Object.keys(functionConfigs).length;
+    return `${numFunctions} functions`;
+  });
+
+  const numVariantsDesc = functionConfigsPromise.then((functionConfigs) => {
+    const numVariants = Object.values(functionConfigs).reduce(
+      (acc, funcConfig) => {
+        return (
+          acc + (funcConfig ? Object.keys(funcConfig.variants || {}).length : 0)
+        );
+      },
+      0,
+    );
+    return `${numVariants} variants`;
+  });
+
+  const numEpisodesDesc = numEpisodesPromise.then(
+    (numEpisodes) => `${numEpisodes.toLocaleString()} episodes`,
+  );
+
+  const numDatasetsDesc = datasetCountsPromise.then(
+    (datasetCounts) => `${datasetCounts.length} datasets`,
+  );
+
+  const numEvaluationRunsDesc = numEvaluationRunsPromise.then(
+    (runs) => `evaluations, ${runs} runs`,
+  );
+
+  // We need to create a special promise for the static evaluations that includes the config count
+  const staticEvaluationsDesc = Promise.all([
+    configPromise,
+    numEvaluationRunsPromise,
+  ]).then(([config, runs]) => {
+    const numEvaluations = Object.keys(config.evaluations || {}).length;
+    return `${numEvaluations} evaluations, ${runs} runs`;
+  });
+
+  const dynamicEvaluationsDesc = Promise.all([
+    numDynamicEvaluationRunProjectsPromise,
+    numDynamicEvaluationRunsPromise,
+  ]).then(([projects, runs]) => `${projects} projects, ${runs} runs`);
 
   return {
-    totalInferences,
-    numFunctions,
-    numVariants,
-    numEpisodes,
-    numDatasets,
-    numEvaluationRuns,
-    numDynamicEvaluationRuns,
-    numDynamicEvaluationRunProjects,
+    totalInferencesDesc,
+    numFunctionsDesc,
+    numVariantsDesc,
+    numEpisodesDesc,
+    numDatasetsDesc,
+    numEvaluationRunsDesc,
+    staticEvaluationsDesc,
+    dynamicEvaluationsDesc,
   };
 }
 
 export default function Home({ loaderData }: Route.ComponentProps) {
   const {
-    totalInferences,
-    numFunctions,
-    numVariants,
-    numEpisodes,
-    numDatasets,
-    numEvaluationRuns,
-    numDynamicEvaluationRuns,
-    numDynamicEvaluationRunProjects,
+    totalInferencesDesc,
+    numFunctionsDesc,
+    numVariantsDesc,
+    numEpisodesDesc,
+    numDatasetsDesc,
+    staticEvaluationsDesc,
+    dynamicEvaluationsDesc,
   } = loaderData;
-  const config = useConfig();
-  const numEvaluations = Object.keys(config.evaluations).length;
 
   return (
     <PageLayout>
@@ -157,19 +202,19 @@ export default function Home({ loaderData }: Route.ComponentProps) {
                 source="/observability/inferences"
                 icon={Inferences}
                 title="Inferences"
-                description={`${totalInferences.toLocaleString()} inferences`}
+                description={totalInferencesDesc}
               />
               <DirectoryCard
                 source="/observability/episodes"
                 icon={Episodes}
                 title="Episodes"
-                description={`${numEpisodes.toLocaleString()} episodes`}
+                description={numEpisodesDesc}
               />
               <DirectoryCard
                 source="/observability/functions"
                 icon={Functions}
                 title="Functions"
-                description={`${numFunctions} functions`}
+                description={numFunctionsDesc}
               />
             </div>
           </div>
@@ -183,7 +228,7 @@ export default function Home({ loaderData }: Route.ComponentProps) {
                 source="/optimization/supervised-fine-tuning"
                 icon={SupervisedFineTuning}
                 title="Supervised Fine-tuning"
-                description={`${numFunctions} functions`}
+                description={numFunctionsDesc}
               />
             </div>
           </div>
@@ -195,25 +240,25 @@ export default function Home({ loaderData }: Route.ComponentProps) {
                 source="/playground"
                 icon={Playground}
                 title="Playground"
-                description={`${numVariants} variants`}
+                description={numVariantsDesc}
               />
               <DirectoryCard
                 source="/datasets"
                 icon={Dataset}
                 title="Datasets"
-                description={`${numDatasets} datasets`}
+                description={numDatasetsDesc}
               />
               <DirectoryCard
                 source="/evaluations"
                 icon={GridCheck}
                 title="Static Evaluations"
-                description={`${numEvaluations} evaluations, ${numEvaluationRuns} runs`}
+                description={staticEvaluationsDesc}
               />
               <DirectoryCard
                 source="/dynamic_evaluations"
                 icon={SequenceChecks}
                 title="Dynamic Evaluations"
-                description={`${numDynamicEvaluationRunProjects} projects, ${numDynamicEvaluationRuns} runs`}
+                description={dynamicEvaluationsDesc}
               />
             </div>
           </div>
