From 73a6412481dd2c25cf17be0759cbd790eb67d2ed Mon Sep 17 00:00:00 2001
From: Nate Kelley <nate@buster.so>
Date: Wed, 8 Oct 2025 17:00:42 -0600
Subject: [PATCH] ensure server side loaders

---
 .../src/api/buster_rest/chats/queryRequests.ts   | 10 ++++++++++
 .../routes/screenshots/chats.$chatId.content.tsx | 16 ++++------------
 .../dashboards.$dashboardId.content.tsx          | 10 ++++++++--
 .../screenshots/metrics.$metricId.content.tsx    | 12 +++++++++---
 4 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/apps/web/src/api/buster_rest/chats/queryRequests.ts b/apps/web/src/api/buster_rest/chats/queryRequests.ts
index 9ef2a5a67..4d06eb209 100644
--- a/apps/web/src/api/buster_rest/chats/queryRequests.ts
+++ b/apps/web/src/api/buster_rest/chats/queryRequests.ts
@@ -153,6 +153,16 @@ export const prefetchGetChat = async (
   return existingData || queryClient.getQueryData(query.queryKey);
 };
 
+export const ensureChatData = async (
+  queryClient: QueryClient,
+  params: Parameters<typeof getChat>[0]
+) => {
+  return await queryClient.ensureQueryData({
+    ...chatQueryKeys.chatsGetChat(params.id),
+    queryFn: () => getChatQueryFn(params, queryClient),
+  });
+};
+
 export const useUpdateChat = (params?: { updateToServer?: boolean }) => {
   const queryClient = useQueryClient();
   const { updateToServer = true } = params || {};
diff --git a/apps/web/src/routes/screenshots/chats.$chatId.content.tsx b/apps/web/src/routes/screenshots/chats.$chatId.content.tsx
index 208826166..1b9d13fb0 100644
--- a/apps/web/src/routes/screenshots/chats.$chatId.content.tsx
+++ b/apps/web/src/routes/screenshots/chats.$chatId.content.tsx
@@ -1,22 +1,14 @@
 import { GetChatScreenshotQuerySchema } from '@buster/server-shared/screenshots';
-import { createFileRoute, redirect } from '@tanstack/react-router';
-import { prefetchGetChat } from '@/api/buster_rest/chats';
-import * as chatLayoutServerContext from '@/context/BusterAssets/chat-server/chatLayoutServer';
+import { createFileRoute } from '@tanstack/react-router';
+import { ensureChatData } from '@/api/buster_rest/chats';
 import { ChatLayout } from '@/layouts/ChatLayout';
 import { DEFAULT_CHAT_ONLY_LAYOUT } from '@/layouts/ChatLayout/config';
 
 export const Route = createFileRoute('/screenshots/chats/$chatId/content')({
-  ...chatLayoutServerContext,
   validateSearch: GetChatScreenshotQuerySchema,
   ssr: true,
-  beforeLoad: async ({ context, params }) => {
-    const chat = await prefetchGetChat({ id: params.chatId }, context.queryClient);
-
-    if (!chat) {
-      throw redirect({
-        to: '/app/home',
-      });
-    }
+  loader: async ({ context, params }) => {
+    await ensureChatData(context.queryClient, { id: params.chatId });
   },
   component: () => {
     return (
diff --git a/apps/web/src/routes/screenshots/dashboards.$dashboardId.content.tsx b/apps/web/src/routes/screenshots/dashboards.$dashboardId.content.tsx
index 443db6294..218d7029a 100644
--- a/apps/web/src/routes/screenshots/dashboards.$dashboardId.content.tsx
+++ b/apps/web/src/routes/screenshots/dashboards.$dashboardId.content.tsx
@@ -9,11 +9,17 @@ export const Route = createFileRoute('/screenshots/dashboards/$dashboardId/conte
   component: RouteComponent,
   ssr: true,
   validateSearch: GetDashboardScreenshotQuerySchema,
-  beforeLoad: async ({ context, params, search }) => {
+  beforeLoad: ({ search }) => {
+    return {
+      version_number: search.version_number,
+    };
+  },
+  loader: async ({ context, params }) => {
+    const { version_number } = context;
     const dashboard = await prefetchGetDashboard({
       queryClient: context.queryClient,
       id: params.dashboardId,
-      version_number: search.version_number,
+      version_number: version_number,
     });
 
     if (!dashboard) {
diff --git a/apps/web/src/routes/screenshots/metrics.$metricId.content.tsx b/apps/web/src/routes/screenshots/metrics.$metricId.content.tsx
index 9c3242659..fbdb64530 100644
--- a/apps/web/src/routes/screenshots/metrics.$metricId.content.tsx
+++ b/apps/web/src/routes/screenshots/metrics.$metricId.content.tsx
@@ -7,15 +7,21 @@ import { MetricViewChartController } from '@/controllers/MetricController/Metric
 export const Route = createFileRoute('/screenshots/metrics/$metricId/content')({
   validateSearch: GetMetricScreenshotQuerySchema,
   ssr: true,
-  beforeLoad: async ({ context, params, search }) => {
+  beforeLoad: ({ search }) => {
+    return {
+      version_number: search.version_number,
+    };
+  },
+  loader: async ({ context, params }) => {
+    const { version_number } = context;
     const [metric, metricData] = await Promise.all([
       prefetchGetMetric(context.queryClient, {
         id: params.metricId,
-        version_number: search.version_number,
+        version_number: version_number,
       }),
       ensureMetricData(context.queryClient, {
         id: params.metricId,
-        version_number: search.version_number,
+        version_number: version_number,
       }),
     ]);
     if (!metric || !metricData) {