ensure server side loaders

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 || {};

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 (

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) {

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) {