Update should take chat screenshot

apps/server/src/api/v2/chats/[id]/GET.ts

@@ -1,12 +1,20 @@
+import { screenshots_task_keys } from '@buster-app/trigger/task-keys';
+import type { TakeChatScreenshotTrigger } from '@buster-app/trigger/task-schemas';
 import { checkPermission } from '@buster/access-controls';
 import type { User } from '@buster/database/queries';
-import { getChatWithDetails, getMessagesForChatWithUserDetails } from '@buster/database/queries';
+import {
+  getChatWithDetails,
+  getMessagesForChatWithUserDetails,
+  getUserOrganizationId,
+} from '@buster/database/queries';
 import {
   GetChatRequestParamsSchema,
   GetChatRequestQuerySchema,
   type GetChatResponse,
 } from '@buster/server-shared/chats';
 import { zValidator } from '@hono/zod-validator';
+import { shouldTakeScreenshot } from '@shared-helpers/screenshots';
+import { tasks } from '@trigger.dev/sdk';
 import { Hono } from 'hono';
 import { HTTPException } from 'hono/http-exception';
 import { throwUnauthorizedError } from '../../../../shared-helpers/asset-public-access';
@@ -40,6 +48,27 @@ const app = new Hono().get(
 
     const response: GetChatResponse = await getChatHandler(getChatHandlerParams);
 
+    const tag = `take-chat-screenshot-${id}`;
+    if (
+      !response.screenshot_taken_at &&
+      (await shouldTakeScreenshot({
+        tag,
+        key: screenshots_task_keys.take_chat_screenshot,
+        context: c,
+      }))
+    ) {
+      tasks.trigger(
+        screenshots_task_keys.take_chat_screenshot,
+        {
+          chatId: id,
+          isNewChatMessage: false,
+          organizationId: (await getUserOrganizationId(user.id))?.organizationId || '',
+          accessToken: c.get('accessToken'),
+        } satisfies TakeChatScreenshotTrigger,
+        { tags: [tag] }
+      );
+    }
+
     return c.json(response);
   }
 );

apps/server/src/api/v2/chats/services/chat-helpers.ts

@@ -207,6 +207,7 @@ export async function buildChatWithMessages(
     permission,
     workspace_sharing: chat.workspaceSharing || 'none',
     workspace_member_count: workspaceMemberCount,
+    screenshot_taken_at: chat.screenshotTakenAt || null,
   };
 }
 

apps/server/src/api/v2/reports/[id]/GET.ts

@@ -107,7 +107,7 @@ const app = new Hono()
             organizationId: (await getUserOrganizationId(user.id))?.organizationId || '',
             accessToken: c.get('accessToken'),
           } satisfies TakeReportScreenshotTrigger,
-          { concurrencyKey: `take-report-screenshot-${reportId}-${versionNumber}` }
+          { tags: [tag] }
         );
       }
 

apps/server/src/shared-helpers/screenshots.ts

@@ -2,6 +2,9 @@ import type { screenshots_task_keys } from '@buster-app/trigger/task-keys';
 import { runs } from '@trigger.dev/sdk';
 import type { Context } from 'hono';
 
+// This helper ensures that we do not run multiple trigger jobs for the same screenshot task concurrently.
+// It checks if a job for the given tag and key is already running or queued before starting a new one.
+
 export const shouldTakeScreenshot = async ({
   tag,
   key,

apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts

@@ -522,11 +522,12 @@ export const analystAgentTask: ReturnType<
         const supabaseUser = await getSupabaseUser(payload.access_token);
         if (supabaseUser) {
           await tasks.trigger(
-            screenshots_task_keys.take_chart_screenshot,
+            screenshots_task_keys.take_chat_screenshot,
             {
               chatId: messageContext.chatId,
               organizationId: messageContext.organizationId,
               accessToken: payload.access_token,
+              isNewChatMessage: true,
             } satisfies TakeChatScreenshotTrigger,
             { concurrencyKey: `take-dashboard-screenshot-${payload.message_id}` }
           );

apps/trigger/src/tasks/screenshots/schemas.ts

@@ -28,6 +28,7 @@ export type TakeReportScreenshotTrigger = z.infer<typeof TakeReportScreenshotTri
 
 export const TakeChatScreenshotTriggerSchema = GetChatScreenshotHandlerArgsSchema.extend({
   chatId: z.string(),
+  isNewChatMessage: z.boolean(),
 });
 
 export type TakeChatScreenshotTrigger = z.infer<typeof TakeChatScreenshotTriggerSchema>;

apps/trigger/src/tasks/screenshots/take-chat-screenshot-handler.ts

@@ -1,22 +1,33 @@
+import { hasChatScreenshotBeenTakenWithin } from '@buster/database/queries';
 import { getChatScreenshot } from '@buster/server-shared/screenshots/methods';
 import { logger, schemaTask } from '@trigger.dev/sdk';
-import { TakeChatScreenshotTriggerSchema } from './schemas';
+import dayjs from 'dayjs';
+import { type TakeChatScreenshotTrigger, TakeChatScreenshotTriggerSchema } from './schemas';
 import { screenshots_task_keys } from './task-keys';
 import { uploadScreenshotHandler } from './upload-screenshot-handler';
 
 export const takeChartScreenshotHandlerTask: ReturnType<
   typeof schemaTask<
-    typeof screenshots_task_keys.take_chart_screenshot,
+    typeof screenshots_task_keys.take_chat_screenshot,
     typeof TakeChatScreenshotTriggerSchema,
     { success: boolean } | undefined
   >
 > = schemaTask({
-  id: screenshots_task_keys.take_chart_screenshot,
+  id: screenshots_task_keys.take_chat_screenshot,
   schema: TakeChatScreenshotTriggerSchema,
   run: async (args) => {
     logger.info('Getting chart screenshot', { args });
 
-    const { chatId, organizationId } = args;
+    const { chatId, isNewChatMessage, organizationId } = args;
+
+    const shouldTakeNewScreenshot = await shouldTakeChatScreenshot({
+      chatId,
+      isNewChatMessage,
+    });
+
+    if (!shouldTakeNewScreenshot) {
+      return;
+    }
 
     const screenshotBuffer = await getChatScreenshot(args);
 
@@ -34,3 +45,18 @@ export const takeChartScreenshotHandlerTask: ReturnType<
     return result;
   },
 });
+
+const shouldTakeChatScreenshot = async (
+  args: Pick<TakeChatScreenshotTrigger, 'chatId' | 'isNewChatMessage'>
+) => {
+  if (args.isNewChatMessage) {
+    return true;
+  }
+
+  const isScreenshotExpired = await hasChatScreenshotBeenTakenWithin(
+    args.chatId,
+    dayjs().subtract(4, 'weeks')
+  );
+
+  return !isScreenshotExpired;
+};

apps/trigger/src/tasks/screenshots/task-keys.ts

@@ -2,6 +2,6 @@ export const screenshots_task_keys = {
   take_metric_screenshot: 'take-metric-screenshot',
   take_dashboard_screenshot: 'take-dashboard-screenshot',
   take_report_screenshot: 'take-report-screenshot',
-  take_chart_screenshot: 'take-chart-screenshot',
+  take_chat_screenshot: 'take-chart-screenshot',
   take_collection_screenshot: 'take-collection-screenshot',
 } as const;

packages/server-shared/src/chats/chat.types.ts

@@ -30,6 +30,7 @@ export const ChatWithMessagesSchema = z.object({
   created_by_id: z.string(),
   created_by_name: z.string(),
   created_by_avatar: z.string().nullable(),
+  screenshot_taken_at: z.string().nullable(),
   ...ShareConfigSchema.shape,
 });