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,
} 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<TakeChatScreenshotTrigger>({
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);
}

View File

@ -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: {

View File

@ -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<TakeDashboardScreenshotTrigger>({
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);
}

View File

@ -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);

View File

@ -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<TakeMetricScreenshotTrigger>({
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);
}

View File

@ -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<TakeReportScreenshotTrigger>({
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);
}

View File

@ -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<string>();
const cooldownCheckingTags = new Set<string>();
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<boolean> => {
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<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 });
}
}