From edcb3ad5777835a4900cc274b58a2c67713f4fbb Mon Sep 17 00:00:00 2001 From: Nate Kelley Date: Wed, 8 Oct 2025 16:24:30 -0600 Subject: [PATCH] Update should take chat screenshot --- apps/server/src/api/v2/chats/[id]/GET.ts | 31 ++++++++++++++++- .../src/api/v2/chats/services/chat-helpers.ts | 1 + apps/server/src/api/v2/reports/[id]/GET.ts | 2 +- apps/server/src/shared-helpers/screenshots.ts | 3 ++ .../analyst-agent-task/analyst-agent-task.ts | 3 +- apps/trigger/src/tasks/screenshots/schemas.ts | 1 + .../take-chat-screenshot-handler.ts | 34 ++++++++++++++++--- .../src/tasks/screenshots/task-keys.ts | 2 +- .../server-shared/src/chats/chat.types.ts | 1 + 9 files changed, 70 insertions(+), 8 deletions(-) diff --git a/apps/server/src/api/v2/chats/[id]/GET.ts b/apps/server/src/api/v2/chats/[id]/GET.ts index a4a1aa6b2..50658ae84 100644 --- a/apps/server/src/api/v2/chats/[id]/GET.ts +++ b/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); } ); diff --git a/apps/server/src/api/v2/chats/services/chat-helpers.ts b/apps/server/src/api/v2/chats/services/chat-helpers.ts index 7aacdbf86..802dbbee2 100644 --- a/apps/server/src/api/v2/chats/services/chat-helpers.ts +++ b/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, }; } diff --git a/apps/server/src/api/v2/reports/[id]/GET.ts b/apps/server/src/api/v2/reports/[id]/GET.ts index ed8415631..d53ca4911 100644 --- a/apps/server/src/api/v2/reports/[id]/GET.ts +++ b/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] } ); } diff --git a/apps/server/src/shared-helpers/screenshots.ts b/apps/server/src/shared-helpers/screenshots.ts index 9759e58f6..44c3dca14 100644 --- a/apps/server/src/shared-helpers/screenshots.ts +++ b/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, diff --git a/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts b/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts index 148de0e0e..27d6b2f43 100644 --- a/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts +++ b/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}` } ); diff --git a/apps/trigger/src/tasks/screenshots/schemas.ts b/apps/trigger/src/tasks/screenshots/schemas.ts index 8e932094a..4ddca1b3d 100644 --- a/apps/trigger/src/tasks/screenshots/schemas.ts +++ b/apps/trigger/src/tasks/screenshots/schemas.ts @@ -28,6 +28,7 @@ export type TakeReportScreenshotTrigger = z.infer; diff --git a/apps/trigger/src/tasks/screenshots/take-chat-screenshot-handler.ts b/apps/trigger/src/tasks/screenshots/take-chat-screenshot-handler.ts index 58837dd66..770647746 100644 --- a/apps/trigger/src/tasks/screenshots/take-chat-screenshot-handler.ts +++ b/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 +) => { + if (args.isNewChatMessage) { + return true; + } + + const isScreenshotExpired = await hasChatScreenshotBeenTakenWithin( + args.chatId, + dayjs().subtract(4, 'weeks') + ); + + return !isScreenshotExpired; +}; diff --git a/apps/trigger/src/tasks/screenshots/task-keys.ts b/apps/trigger/src/tasks/screenshots/task-keys.ts index 29780890d..ae60bade8 100644 --- a/apps/trigger/src/tasks/screenshots/task-keys.ts +++ b/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; diff --git a/packages/server-shared/src/chats/chat.types.ts b/packages/server-shared/src/chats/chat.types.ts index 8346b97de..6e3ff014c 100644 --- a/packages/server-shared/src/chats/chat.types.ts +++ b/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, });