From 03ed4145995b5ebb6b85780ad3f5e8804370ae7f Mon Sep 17 00:00:00 2001 From: Nate Kelley Date: Thu, 9 Oct 2025 08:21:06 -0600 Subject: [PATCH] shared screenshot handler --- apps/server/src/api/v2/chats/[id]/GET.ts | 35 ++++----- .../src/api/v2/dashboards/[id]/GET.test.ts | 76 +++++-------------- apps/server/src/api/v2/dashboards/[id]/GET.ts | 33 +++----- .../src/api/v2/dashboards/[id]/sharing/PUT.ts | 12 +-- .../src/api/v2/metric_files/[id]/data/GET.ts | 36 ++++----- apps/server/src/api/v2/reports/[id]/GET.ts | 32 +++----- apps/server/src/shared-helpers/screenshots.ts | 41 ++++++++-- 7 files changed, 104 insertions(+), 161 deletions(-) diff --git a/apps/server/src/api/v2/chats/[id]/GET.ts b/apps/server/src/api/v2/chats/[id]/GET.ts index 88c7d00c6..34fb87946 100644 --- a/apps/server/src/api/v2/chats/[id]/GET.ts +++ b/apps/server/src/api/v2/chats/[id]/GET.ts @@ -13,8 +13,7 @@ import { 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 { triggerScreenshotIfNeeded } from '@shared-helpers/screenshots'; import { Hono } from 'hono'; import { HTTPException } from 'hono/http-exception'; import { throwUnauthorizedError } from '../../../../shared-helpers/asset-public-access'; @@ -48,26 +47,18 @@ 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], idempotencyKey: tag } - ); - } + await triggerScreenshotIfNeeded({ + tag: `take-chat-screenshot-${id}`, + key: screenshots_task_keys.take_chat_screenshot, + context: c, + payload: { + chatId: id, + isNewChatMessage: false, + organizationId: (await getUserOrganizationId(user.id))?.organizationId || '', + accessToken: c.get('accessToken'), + }, + shouldTrigger: !response.screenshot_taken_at, + }); return c.json(response); } diff --git a/apps/server/src/api/v2/dashboards/[id]/GET.test.ts b/apps/server/src/api/v2/dashboards/[id]/GET.test.ts index b4f184e3c..06581b772 100644 --- a/apps/server/src/api/v2/dashboards/[id]/GET.test.ts +++ b/apps/server/src/api/v2/dashboards/[id]/GET.test.ts @@ -176,11 +176,7 @@ describe('getDashboardHandler', () => { describe('successful requests', () => { it('should return dashboard data for valid dashboard ID', async () => { - const result = await getDashboardHandler( - { dashboardId: 'dashboard-123' }, - mockUser, - mockContext - ); + const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser); expect(result.dashboard.id).toBe('dashboard-123'); expect(result.dashboard.name).toBe('Test Dashboard'); @@ -223,8 +219,7 @@ describe('getDashboardHandler', () => { const result = await getDashboardHandler( { dashboardId: 'dashboard-123', versionNumber: 1 }, - mockUser, - mockContext + mockUser ); expect((result.dashboard.config as any).name).toBe('Version 1'); @@ -245,8 +240,7 @@ describe('getDashboardHandler', () => { const result = await getDashboardHandler( { dashboardId: 'dashboard-123', password: 'secret123' }, - mockUser, - mockContext + mockUser ); expect(result.permission).toBe('can_view'); @@ -258,7 +252,7 @@ describe('getDashboardHandler', () => { mockGetDashboardById.mockResolvedValue(null); await expect( - getDashboardHandler({ dashboardId: 'nonexistent-dashboard' }, mockUser, mockContext) + getDashboardHandler({ dashboardId: 'nonexistent-dashboard' }, mockUser) ).rejects.toThrow(new HTTPException(404, { message: 'Dashboard not found' })); }); @@ -273,9 +267,7 @@ describe('getDashboardHandler', () => { effectiveRole: undefined, }); - await expect( - getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser, mockContext) - ).rejects.toThrow( + await expect(getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser)).rejects.toThrow( new HTTPException(403, { message: "You don't have permission to view this dashboard" }) ); }); @@ -295,9 +287,7 @@ describe('getDashboardHandler', () => { effectiveRole: undefined, }); - await expect( - getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser, mockContext) - ).rejects.toThrow( + await expect(getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser)).rejects.toThrow( new HTTPException(403, { message: 'Public access to this dashboard has expired' }) ); }); @@ -314,9 +304,9 @@ describe('getDashboardHandler', () => { effectiveRole: undefined, }); - await expect( - getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser, mockContext) - ).rejects.toThrow(new HTTPException(418, { message: 'Password required for public access' })); + await expect(getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser)).rejects.toThrow( + new HTTPException(418, { message: 'Password required for public access' }) + ); }); it('should throw 403 when incorrect password is provided', async () => { @@ -332,11 +322,7 @@ describe('getDashboardHandler', () => { }); await expect( - getDashboardHandler( - { dashboardId: 'dashboard-123', password: 'wrong-password' }, - mockUser, - mockContext - ) + getDashboardHandler({ dashboardId: 'dashboard-123', password: 'wrong-password' }, mockUser) ).rejects.toThrow( new HTTPException(403, { message: 'Incorrect password for public access' }) ); @@ -344,11 +330,7 @@ describe('getDashboardHandler', () => { it('should throw 404 when requested version does not exist', async () => { await expect( - getDashboardHandler( - { dashboardId: 'dashboard-123', versionNumber: 99 }, - mockUser, - mockContext - ) + getDashboardHandler({ dashboardId: 'dashboard-123', versionNumber: 99 }, mockUser) ).rejects.toThrow(new HTTPException(404, { message: 'Version 99 not found' })); }); @@ -366,11 +348,7 @@ describe('getDashboardHandler', () => { mockGetDashboardById.mockResolvedValue(versionedDashboard); await expect( - getDashboardHandler( - { dashboardId: 'dashboard-123', versionNumber: 1 }, - mockUser, - mockContext - ) + getDashboardHandler({ dashboardId: 'dashboard-123', versionNumber: 1 }, mockUser) ).rejects.toThrow(new HTTPException(404, { message: 'Version 1 not found' })); }); }); @@ -399,11 +377,7 @@ describe('getDashboardHandler', () => { }; mockGetDashboardById.mockResolvedValue(versionedDashboard); - const result = await getDashboardHandler( - { dashboardId: 'dashboard-123' }, - mockUser, - mockContext - ); + const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser); expect(result.versions).toEqual([ { version_number: 1, updated_at: '2023-01-01T00:00:00Z' }, @@ -430,11 +404,7 @@ describe('getDashboardHandler', () => { }; mockGetDashboardById.mockResolvedValue(versionedDashboard); - const result = await getDashboardHandler( - { dashboardId: 'dashboard-123' }, - mockUser, - mockContext - ); + const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser); expect(result.dashboard.version_number).toBe(2); expect(result.dashboard.config).toEqual(mockDashboardContent); // From current content @@ -448,11 +418,7 @@ describe('getDashboardHandler', () => { effectiveRole: 'can_edit', }); - const result = await getDashboardHandler( - { dashboardId: 'dashboard-123' }, - mockUser, - mockContext - ); + const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser); expect(result.permission).toBe('can_edit'); }); @@ -468,11 +434,7 @@ describe('getDashboardHandler', () => { effectiveRole: undefined, }); - const result = await getDashboardHandler( - { dashboardId: 'dashboard-123' }, - mockUser, - mockContext - ); + const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser); expect(result.permission).toBe('can_view'); }); @@ -480,11 +442,7 @@ describe('getDashboardHandler', () => { describe('response structure', () => { it('should return properly formatted dashboard data', async () => { - const result = await getDashboardHandler( - { dashboardId: 'dashboard-123' }, - mockUser, - mockContext - ); + const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser); expect(result).toMatchObject({ dashboard: { diff --git a/apps/server/src/api/v2/dashboards/[id]/GET.ts b/apps/server/src/api/v2/dashboards/[id]/GET.ts index 57c7c9b68..9be7778cb 100644 --- a/apps/server/src/api/v2/dashboards/[id]/GET.ts +++ b/apps/server/src/api/v2/dashboards/[id]/GET.ts @@ -17,8 +17,7 @@ import { import type { DashboardYml } from '@buster/server-shared/dashboards'; import type { VerificationStatus } from '@buster/server-shared/share'; import { zValidator } from '@hono/zod-validator'; -import { shouldTakeScreenshot } from '@shared-helpers/screenshots'; -import { tasks } from '@trigger.dev/sdk'; +import { triggerScreenshotIfNeeded } from '@shared-helpers/screenshots'; import { type Context, Hono } from 'hono'; import { HTTPException } from 'hono/http-exception'; import yaml from 'js-yaml'; @@ -54,25 +53,17 @@ const app = new Hono().get( user ); - const tag = `take-dashboard-screenshot-${id}`; - if ( - await shouldTakeScreenshot({ - tag, - key: screenshots_task_keys.take_dashboard_screenshot, - context: c, - }) - ) { - tasks.trigger( - screenshots_task_keys.take_dashboard_screenshot, - { - dashboardId: id, - organizationId: (await getUserOrganizationId(user.id))?.organizationId || '', - accessToken: c.get('accessToken'), - isOnSaveEvent: false, - } satisfies TakeDashboardScreenshotTrigger, - { tags: [tag], idempotencyKey: tag } - ); - } + await triggerScreenshotIfNeeded({ + tag: `take-dashboard-screenshot-${id}`, + key: screenshots_task_keys.take_dashboard_screenshot, + context: c, + payload: { + dashboardId: id, + organizationId: (await getUserOrganizationId(user.id))?.organizationId || '', + accessToken: c.get('accessToken'), + isOnSaveEvent: false, + }, + }); return c.json(response); } diff --git a/apps/server/src/api/v2/dashboards/[id]/sharing/PUT.ts b/apps/server/src/api/v2/dashboards/[id]/sharing/PUT.ts index 96fb95c5e..3109bf26b 100644 --- a/apps/server/src/api/v2/dashboards/[id]/sharing/PUT.ts +++ b/apps/server/src/api/v2/dashboards/[id]/sharing/PUT.ts @@ -17,8 +17,7 @@ import { getDashboardHandler } from '../GET'; export async function updateDashboardShareHandler( dashboardId: string, request: ShareUpdateRequest, - user: User & { organizationId: string }, - c: Context + user: User & { organizationId: string } ) { // Check if dashboard exists const dashboard = await getDashboardById({ dashboardId }); @@ -105,11 +104,7 @@ export async function updateDashboardShareHandler( workspace_sharing, }); - const updatedDashboard: GetDashboardResponse = await getDashboardHandler( - { dashboardId }, - user, - c - ); + const updatedDashboard: GetDashboardResponse = await getDashboardHandler({ dashboardId }, user); return updatedDashboard; } @@ -135,8 +130,7 @@ const app = new Hono().put('/', zValidator('json', ShareUpdateRequestSchema), as { ...user, organizationId: userOrg.organizationId, - }, - c + } ); return c.json(updatedDashboard); diff --git a/apps/server/src/api/v2/metric_files/[id]/data/GET.ts b/apps/server/src/api/v2/metric_files/[id]/data/GET.ts index a926db947..875f98c10 100644 --- a/apps/server/src/api/v2/metric_files/[id]/data/GET.ts +++ b/apps/server/src/api/v2/metric_files/[id]/data/GET.ts @@ -3,8 +3,7 @@ import type { TakeMetricScreenshotTrigger } from '@buster-app/trigger/task-schem import { getUserOrganizationId } from '@buster/database/queries'; import { MetricDataParamsSchema, MetricDataQuerySchema } from '@buster/server-shared'; import { zValidator } from '@hono/zod-validator'; -import { shouldTakeScreenshot } from '@shared-helpers/screenshots'; -import { runs, tasks } from '@trigger.dev/sdk'; +import { triggerScreenshotIfNeeded } from '@shared-helpers/screenshots'; import dayjs from 'dayjs'; import { Hono } from 'hono'; import { standardErrorHandler } from '../../../../../utils/response'; @@ -29,27 +28,18 @@ const app = new Hono() password ); - const tag = `take-metric-screenshot-${id}`; - if ( - await shouldTakeScreenshot({ - tag, - key: screenshots_task_keys.take_metric_screenshot, - context: c, - }) - ) { - const organizationId = - (await getUserOrganizationId(user.id).then((res) => res?.organizationId)) || ''; - tasks.trigger( - screenshots_task_keys.take_metric_screenshot, - { - metricId: id, - isOnSaveEvent: false, - accessToken: c.get('accessToken'), - organizationId, - } satisfies TakeMetricScreenshotTrigger, - { tags: [tag], idempotencyKey: tag } - ); - } + await triggerScreenshotIfNeeded({ + tag: `take-metric-screenshot-${id}`, + key: screenshots_task_keys.take_metric_screenshot, + context: c, + payload: { + metricId: id, + isOnSaveEvent: false, + accessToken: c.get('accessToken'), + organizationId: + (await getUserOrganizationId(user.id).then((res) => res?.organizationId)) || '', + }, + }); return c.json(response); } diff --git a/apps/server/src/api/v2/reports/[id]/GET.ts b/apps/server/src/api/v2/reports/[id]/GET.ts index dbccb4f26..d28cc1dee 100644 --- a/apps/server/src/api/v2/reports/[id]/GET.ts +++ b/apps/server/src/api/v2/reports/[id]/GET.ts @@ -13,8 +13,7 @@ import { type GetReportResponse, } from '@buster/server-shared/reports'; import { zValidator } from '@hono/zod-validator'; -import { shouldTakeScreenshot } from '@shared-helpers/screenshots'; -import { tasks } from '@trigger.dev/sdk'; +import { triggerScreenshotIfNeeded } from '@shared-helpers/screenshots'; import { Hono } from 'hono'; import { HTTPException } from 'hono/http-exception'; import { throwUnauthorizedError } from '../../../../shared-helpers/asset-public-access'; @@ -91,25 +90,16 @@ const app = new Hono() password ); - const tag = `take-report-screenshot-${reportId}`; - - if ( - await shouldTakeScreenshot({ - tag, - key: screenshots_task_keys.take_report_screenshot, - context: c, - }) - ) { - tasks.trigger( - screenshots_task_keys.take_report_screenshot, - { - reportId, - organizationId: (await getUserOrganizationId(user.id))?.organizationId || '', - accessToken: c.get('accessToken'), - } satisfies TakeReportScreenshotTrigger, - { tags: [tag], idempotencyKey: tag } - ); - } + await triggerScreenshotIfNeeded({ + tag: `take-report-screenshot-${reportId}`, + key: screenshots_task_keys.take_report_screenshot, + context: c, + payload: { + reportId, + organizationId: (await getUserOrganizationId(user.id))?.organizationId || '', + accessToken: c.get('accessToken'), + }, + }); return c.json(response); } diff --git a/apps/server/src/shared-helpers/screenshots.ts b/apps/server/src/shared-helpers/screenshots.ts index fb61a5b60..02c477419 100644 --- a/apps/server/src/shared-helpers/screenshots.ts +++ b/apps/server/src/shared-helpers/screenshots.ts @@ -1,14 +1,14 @@ import type { screenshots_task_keys } from '@buster-app/trigger/task-keys'; -import { runs } from '@trigger.dev/sdk'; +import { runs, tasks } 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. -const currentlyCheckingTags = new Set(); +const cooldownCheckingTags = new Set(); const CACHE_TAG_EXPIRATION_TIME = 1000 * 30; // 30 seconds -export const shouldTakeScreenshot = async ({ +const shouldTakeScreenshot = async ({ tag, key, context, @@ -19,11 +19,11 @@ export const shouldTakeScreenshot = async ({ }): Promise => { const hasReferrer = !!context.req.header('referer'); - if (!hasReferrer || currentlyCheckingTags.has(tag)) { + if (!hasReferrer || cooldownCheckingTags.has(tag)) { return false; } - currentlyCheckingTags.add(tag); + cooldownCheckingTags.add(tag); try { const lastTask = await runs @@ -41,7 +41,36 @@ export const shouldTakeScreenshot = async ({ throw error; } finally { setTimeout(() => { - currentlyCheckingTags.delete(tag); + cooldownCheckingTags.delete(tag); }, CACHE_TAG_EXPIRATION_TIME); } }; + +/** + * Generic handler to conditionally trigger screenshot tasks + * @template TTrigger - The trigger payload type for the screenshot task + */ +export async function triggerScreenshotIfNeeded({ + tag, + key, + context, + payload, + shouldTrigger = true, +}: { + tag: string; + key: (typeof screenshots_task_keys)[keyof typeof screenshots_task_keys]; + context: Context; + payload: TTrigger; + shouldTrigger?: boolean; +}): Promise { + if ( + shouldTrigger && + (await shouldTakeScreenshot({ + tag, + key, + context, + })) + ) { + tasks.trigger(key, payload, { tags: [tag], idempotencyKey: tag }); + } +}