shared screenshot handler

This commit is contained in:
Nate Kelley 2025-10-09 08:21:06 -06:00
parent 08bb8c04f3
commit 03ed414599
No known key found for this signature in database
GPG Key ID: FD90372AB8D98B4F
7 changed files with 104 additions and 161 deletions

View File

@ -13,8 +13,7 @@ import {
type GetChatResponse, type GetChatResponse,
} from '@buster/server-shared/chats'; } from '@buster/server-shared/chats';
import { zValidator } from '@hono/zod-validator'; import { zValidator } from '@hono/zod-validator';
import { shouldTakeScreenshot } from '@shared-helpers/screenshots'; import { triggerScreenshotIfNeeded } from '@shared-helpers/screenshots';
import { tasks } from '@trigger.dev/sdk';
import { Hono } from 'hono'; import { Hono } from 'hono';
import { HTTPException } from 'hono/http-exception'; import { HTTPException } from 'hono/http-exception';
import { throwUnauthorizedError } from '../../../../shared-helpers/asset-public-access'; import { throwUnauthorizedError } from '../../../../shared-helpers/asset-public-access';
@ -48,26 +47,18 @@ const app = new Hono().get(
const response: GetChatResponse = await getChatHandler(getChatHandlerParams); const response: GetChatResponse = await getChatHandler(getChatHandlerParams);
const tag = `take-chat-screenshot-${id}`; await triggerScreenshotIfNeeded<TakeChatScreenshotTrigger>({
if ( tag: `take-chat-screenshot-${id}`,
!response.screenshot_taken_at && key: screenshots_task_keys.take_chat_screenshot,
(await shouldTakeScreenshot({ context: c,
tag, payload: {
key: screenshots_task_keys.take_chat_screenshot, chatId: id,
context: c, isNewChatMessage: false,
})) organizationId: (await getUserOrganizationId(user.id))?.organizationId || '',
) { accessToken: c.get('accessToken'),
tasks.trigger( },
screenshots_task_keys.take_chat_screenshot, shouldTrigger: !response.screenshot_taken_at,
{ });
chatId: id,
isNewChatMessage: false,
organizationId: (await getUserOrganizationId(user.id))?.organizationId || '',
accessToken: c.get('accessToken'),
} satisfies TakeChatScreenshotTrigger,
{ tags: [tag], idempotencyKey: tag }
);
}
return c.json(response); return c.json(response);
} }

View File

@ -176,11 +176,7 @@ describe('getDashboardHandler', () => {
describe('successful requests', () => { describe('successful requests', () => {
it('should return dashboard data for valid dashboard ID', async () => { it('should return dashboard data for valid dashboard ID', async () => {
const result = await getDashboardHandler( const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser);
{ dashboardId: 'dashboard-123' },
mockUser,
mockContext
);
expect(result.dashboard.id).toBe('dashboard-123'); expect(result.dashboard.id).toBe('dashboard-123');
expect(result.dashboard.name).toBe('Test Dashboard'); expect(result.dashboard.name).toBe('Test Dashboard');
@ -223,8 +219,7 @@ describe('getDashboardHandler', () => {
const result = await getDashboardHandler( const result = await getDashboardHandler(
{ dashboardId: 'dashboard-123', versionNumber: 1 }, { dashboardId: 'dashboard-123', versionNumber: 1 },
mockUser, mockUser
mockContext
); );
expect((result.dashboard.config as any).name).toBe('Version 1'); expect((result.dashboard.config as any).name).toBe('Version 1');
@ -245,8 +240,7 @@ describe('getDashboardHandler', () => {
const result = await getDashboardHandler( const result = await getDashboardHandler(
{ dashboardId: 'dashboard-123', password: 'secret123' }, { dashboardId: 'dashboard-123', password: 'secret123' },
mockUser, mockUser
mockContext
); );
expect(result.permission).toBe('can_view'); expect(result.permission).toBe('can_view');
@ -258,7 +252,7 @@ describe('getDashboardHandler', () => {
mockGetDashboardById.mockResolvedValue(null); mockGetDashboardById.mockResolvedValue(null);
await expect( await expect(
getDashboardHandler({ dashboardId: 'nonexistent-dashboard' }, mockUser, mockContext) getDashboardHandler({ dashboardId: 'nonexistent-dashboard' }, mockUser)
).rejects.toThrow(new HTTPException(404, { message: 'Dashboard not found' })); ).rejects.toThrow(new HTTPException(404, { message: 'Dashboard not found' }));
}); });
@ -273,9 +267,7 @@ describe('getDashboardHandler', () => {
effectiveRole: undefined, effectiveRole: undefined,
}); });
await expect( await expect(getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser)).rejects.toThrow(
getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser, mockContext)
).rejects.toThrow(
new HTTPException(403, { message: "You don't have permission to view this dashboard" }) new HTTPException(403, { message: "You don't have permission to view this dashboard" })
); );
}); });
@ -295,9 +287,7 @@ describe('getDashboardHandler', () => {
effectiveRole: undefined, effectiveRole: undefined,
}); });
await expect( await expect(getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser)).rejects.toThrow(
getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser, mockContext)
).rejects.toThrow(
new HTTPException(403, { message: 'Public access to this dashboard has expired' }) new HTTPException(403, { message: 'Public access to this dashboard has expired' })
); );
}); });
@ -314,9 +304,9 @@ describe('getDashboardHandler', () => {
effectiveRole: undefined, effectiveRole: undefined,
}); });
await expect( await expect(getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser)).rejects.toThrow(
getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser, mockContext) new HTTPException(418, { message: 'Password required for public access' })
).rejects.toThrow(new HTTPException(418, { message: 'Password required for public access' })); );
}); });
it('should throw 403 when incorrect password is provided', async () => { it('should throw 403 when incorrect password is provided', async () => {
@ -332,11 +322,7 @@ describe('getDashboardHandler', () => {
}); });
await expect( await expect(
getDashboardHandler( getDashboardHandler({ dashboardId: 'dashboard-123', password: 'wrong-password' }, mockUser)
{ dashboardId: 'dashboard-123', password: 'wrong-password' },
mockUser,
mockContext
)
).rejects.toThrow( ).rejects.toThrow(
new HTTPException(403, { message: 'Incorrect password for public access' }) 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 () => { it('should throw 404 when requested version does not exist', async () => {
await expect( await expect(
getDashboardHandler( getDashboardHandler({ dashboardId: 'dashboard-123', versionNumber: 99 }, mockUser)
{ dashboardId: 'dashboard-123', versionNumber: 99 },
mockUser,
mockContext
)
).rejects.toThrow(new HTTPException(404, { message: 'Version 99 not found' })); ).rejects.toThrow(new HTTPException(404, { message: 'Version 99 not found' }));
}); });
@ -366,11 +348,7 @@ describe('getDashboardHandler', () => {
mockGetDashboardById.mockResolvedValue(versionedDashboard); mockGetDashboardById.mockResolvedValue(versionedDashboard);
await expect( await expect(
getDashboardHandler( getDashboardHandler({ dashboardId: 'dashboard-123', versionNumber: 1 }, mockUser)
{ dashboardId: 'dashboard-123', versionNumber: 1 },
mockUser,
mockContext
)
).rejects.toThrow(new HTTPException(404, { message: 'Version 1 not found' })); ).rejects.toThrow(new HTTPException(404, { message: 'Version 1 not found' }));
}); });
}); });
@ -399,11 +377,7 @@ describe('getDashboardHandler', () => {
}; };
mockGetDashboardById.mockResolvedValue(versionedDashboard); mockGetDashboardById.mockResolvedValue(versionedDashboard);
const result = await getDashboardHandler( const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser);
{ dashboardId: 'dashboard-123' },
mockUser,
mockContext
);
expect(result.versions).toEqual([ expect(result.versions).toEqual([
{ version_number: 1, updated_at: '2023-01-01T00:00:00Z' }, { version_number: 1, updated_at: '2023-01-01T00:00:00Z' },
@ -430,11 +404,7 @@ describe('getDashboardHandler', () => {
}; };
mockGetDashboardById.mockResolvedValue(versionedDashboard); mockGetDashboardById.mockResolvedValue(versionedDashboard);
const result = await getDashboardHandler( const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser);
{ dashboardId: 'dashboard-123' },
mockUser,
mockContext
);
expect(result.dashboard.version_number).toBe(2); expect(result.dashboard.version_number).toBe(2);
expect(result.dashboard.config).toEqual(mockDashboardContent); // From current content expect(result.dashboard.config).toEqual(mockDashboardContent); // From current content
@ -448,11 +418,7 @@ describe('getDashboardHandler', () => {
effectiveRole: 'can_edit', effectiveRole: 'can_edit',
}); });
const result = await getDashboardHandler( const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser);
{ dashboardId: 'dashboard-123' },
mockUser,
mockContext
);
expect(result.permission).toBe('can_edit'); expect(result.permission).toBe('can_edit');
}); });
@ -468,11 +434,7 @@ describe('getDashboardHandler', () => {
effectiveRole: undefined, effectiveRole: undefined,
}); });
const result = await getDashboardHandler( const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser);
{ dashboardId: 'dashboard-123' },
mockUser,
mockContext
);
expect(result.permission).toBe('can_view'); expect(result.permission).toBe('can_view');
}); });
@ -480,11 +442,7 @@ describe('getDashboardHandler', () => {
describe('response structure', () => { describe('response structure', () => {
it('should return properly formatted dashboard data', async () => { it('should return properly formatted dashboard data', async () => {
const result = await getDashboardHandler( const result = await getDashboardHandler({ dashboardId: 'dashboard-123' }, mockUser);
{ dashboardId: 'dashboard-123' },
mockUser,
mockContext
);
expect(result).toMatchObject({ expect(result).toMatchObject({
dashboard: { dashboard: {

View File

@ -17,8 +17,7 @@ import {
import type { DashboardYml } from '@buster/server-shared/dashboards'; import type { DashboardYml } from '@buster/server-shared/dashboards';
import type { VerificationStatus } from '@buster/server-shared/share'; import type { VerificationStatus } from '@buster/server-shared/share';
import { zValidator } from '@hono/zod-validator'; import { zValidator } from '@hono/zod-validator';
import { shouldTakeScreenshot } from '@shared-helpers/screenshots'; import { triggerScreenshotIfNeeded } from '@shared-helpers/screenshots';
import { tasks } from '@trigger.dev/sdk';
import { type Context, Hono } from 'hono'; import { type Context, Hono } from 'hono';
import { HTTPException } from 'hono/http-exception'; import { HTTPException } from 'hono/http-exception';
import yaml from 'js-yaml'; import yaml from 'js-yaml';
@ -54,25 +53,17 @@ const app = new Hono().get(
user user
); );
const tag = `take-dashboard-screenshot-${id}`; await triggerScreenshotIfNeeded<TakeDashboardScreenshotTrigger>({
if ( tag: `take-dashboard-screenshot-${id}`,
await shouldTakeScreenshot({ key: screenshots_task_keys.take_dashboard_screenshot,
tag, context: c,
key: screenshots_task_keys.take_dashboard_screenshot, payload: {
context: c, dashboardId: id,
}) organizationId: (await getUserOrganizationId(user.id))?.organizationId || '',
) { accessToken: c.get('accessToken'),
tasks.trigger( isOnSaveEvent: false,
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 }
);
}
return c.json(response); return c.json(response);
} }

View File

@ -17,8 +17,7 @@ import { getDashboardHandler } from '../GET';
export async function updateDashboardShareHandler( export async function updateDashboardShareHandler(
dashboardId: string, dashboardId: string,
request: ShareUpdateRequest, request: ShareUpdateRequest,
user: User & { organizationId: string }, user: User & { organizationId: string }
c: Context
) { ) {
// Check if dashboard exists // Check if dashboard exists
const dashboard = await getDashboardById({ dashboardId }); const dashboard = await getDashboardById({ dashboardId });
@ -105,11 +104,7 @@ export async function updateDashboardShareHandler(
workspace_sharing, workspace_sharing,
}); });
const updatedDashboard: GetDashboardResponse = await getDashboardHandler( const updatedDashboard: GetDashboardResponse = await getDashboardHandler({ dashboardId }, user);
{ dashboardId },
user,
c
);
return updatedDashboard; return updatedDashboard;
} }
@ -135,8 +130,7 @@ const app = new Hono().put('/', zValidator('json', ShareUpdateRequestSchema), as
{ {
...user, ...user,
organizationId: userOrg.organizationId, organizationId: userOrg.organizationId,
}, }
c
); );
return c.json(updatedDashboard); return c.json(updatedDashboard);

View File

@ -3,8 +3,7 @@ import type { TakeMetricScreenshotTrigger } from '@buster-app/trigger/task-schem
import { getUserOrganizationId } from '@buster/database/queries'; import { getUserOrganizationId } from '@buster/database/queries';
import { MetricDataParamsSchema, MetricDataQuerySchema } from '@buster/server-shared'; import { MetricDataParamsSchema, MetricDataQuerySchema } from '@buster/server-shared';
import { zValidator } from '@hono/zod-validator'; import { zValidator } from '@hono/zod-validator';
import { shouldTakeScreenshot } from '@shared-helpers/screenshots'; import { triggerScreenshotIfNeeded } from '@shared-helpers/screenshots';
import { runs, tasks } from '@trigger.dev/sdk';
import dayjs from 'dayjs'; import dayjs from 'dayjs';
import { Hono } from 'hono'; import { Hono } from 'hono';
import { standardErrorHandler } from '../../../../../utils/response'; import { standardErrorHandler } from '../../../../../utils/response';
@ -29,27 +28,18 @@ const app = new Hono()
password password
); );
const tag = `take-metric-screenshot-${id}`; await triggerScreenshotIfNeeded<TakeMetricScreenshotTrigger>({
if ( tag: `take-metric-screenshot-${id}`,
await shouldTakeScreenshot({ key: screenshots_task_keys.take_metric_screenshot,
tag, context: c,
key: screenshots_task_keys.take_metric_screenshot, payload: {
context: c, metricId: id,
}) isOnSaveEvent: false,
) { accessToken: c.get('accessToken'),
const organizationId = organizationId:
(await getUserOrganizationId(user.id).then((res) => res?.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 }
);
}
return c.json(response); return c.json(response);
} }

View File

@ -13,8 +13,7 @@ import {
type GetReportResponse, type GetReportResponse,
} from '@buster/server-shared/reports'; } from '@buster/server-shared/reports';
import { zValidator } from '@hono/zod-validator'; import { zValidator } from '@hono/zod-validator';
import { shouldTakeScreenshot } from '@shared-helpers/screenshots'; import { triggerScreenshotIfNeeded } from '@shared-helpers/screenshots';
import { tasks } from '@trigger.dev/sdk';
import { Hono } from 'hono'; import { Hono } from 'hono';
import { HTTPException } from 'hono/http-exception'; import { HTTPException } from 'hono/http-exception';
import { throwUnauthorizedError } from '../../../../shared-helpers/asset-public-access'; import { throwUnauthorizedError } from '../../../../shared-helpers/asset-public-access';
@ -91,25 +90,16 @@ const app = new Hono()
password password
); );
const tag = `take-report-screenshot-${reportId}`; await triggerScreenshotIfNeeded<TakeReportScreenshotTrigger>({
tag: `take-report-screenshot-${reportId}`,
if ( key: screenshots_task_keys.take_report_screenshot,
await shouldTakeScreenshot({ context: c,
tag, payload: {
key: screenshots_task_keys.take_report_screenshot, reportId,
context: c, organizationId: (await getUserOrganizationId(user.id))?.organizationId || '',
}) accessToken: c.get('accessToken'),
) { },
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 }
);
}
return c.json(response); return c.json(response);
} }

View File

@ -1,14 +1,14 @@
import type { screenshots_task_keys } from '@buster-app/trigger/task-keys'; 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'; import type { Context } from 'hono';
// This helper ensures that we do not run multiple trigger jobs for the same screenshot task concurrently. // 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. // 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<string>(); const cooldownCheckingTags = new Set<string>();
const CACHE_TAG_EXPIRATION_TIME = 1000 * 30; // 30 seconds const CACHE_TAG_EXPIRATION_TIME = 1000 * 30; // 30 seconds
export const shouldTakeScreenshot = async ({ const shouldTakeScreenshot = async ({
tag, tag,
key, key,
context, context,
@ -19,11 +19,11 @@ export const shouldTakeScreenshot = async ({
}): Promise<boolean> => { }): Promise<boolean> => {
const hasReferrer = !!context.req.header('referer'); const hasReferrer = !!context.req.header('referer');
if (!hasReferrer || currentlyCheckingTags.has(tag)) { if (!hasReferrer || cooldownCheckingTags.has(tag)) {
return false; return false;
} }
currentlyCheckingTags.add(tag); cooldownCheckingTags.add(tag);
try { try {
const lastTask = await runs const lastTask = await runs
@ -41,7 +41,36 @@ export const shouldTakeScreenshot = async ({
throw error; throw error;
} finally { } finally {
setTimeout(() => { setTimeout(() => {
currentlyCheckingTags.delete(tag); cooldownCheckingTags.delete(tag);
}, CACHE_TAG_EXPIRATION_TIME); }, 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<TTrigger>({
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<void> {
if (
shouldTrigger &&
(await shouldTakeScreenshot({
tag,
key,
context,
}))
) {
tasks.trigger(key, payload, { tags: [tag], idempotencyKey: tag });
}
}