From f15207591790261ff7a5c0719231703f89e9c3b4 Mon Sep 17 00:00:00 2001
From: Nate Kelley <nate@buster.so>
Date: Wed, 8 Oct 2025 15:50:31 -0600
Subject: [PATCH] should take screenshot logic

---
 .../src/api/v2/metric_files/[id]/data/GET.ts  | 19 +++++-------
 apps/server/src/shared-helpers/screenshots.ts | 30 +++++++++++++++++++
 .../src/api/buster_rest/metrics/requests.ts   |  5 ++--
 .../src/metrics/requests.types.ts             |  1 -
 .../src/metrics/responses.types.ts            |  1 -
 5 files changed, 41 insertions(+), 15 deletions(-)
 create mode 100644 apps/server/src/shared-helpers/screenshots.ts

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 aa23060d0..092977feb 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,9 +3,10 @@ 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 dayjs from 'dayjs';
 import { Hono } from 'hono';
-import { id } from 'zod/v4/locales';
 import { standardErrorHandler } from '../../../../../utils/response';
 import { getMetricDataHandler } from './get-metric-data';
 
@@ -16,8 +17,7 @@ const app = new Hono()
     zValidator('query', MetricDataQuerySchema),
     async (c) => {
       const { id } = c.req.valid('param');
-      const { limit, version_number, report_file_id, password, is_screenshot } =
-        c.req.valid('query');
+      const { limit, version_number, report_file_id, password } = c.req.valid('query');
       const user = c.get('busterUser');
 
       const response = await getMetricDataHandler(
@@ -30,16 +30,13 @@ const app = new Hono()
       );
 
       const tag = `take-metric-screenshot-${id}-${version_number}`;
-      const lastTask = await runs
-        .list({
-          status: ['EXECUTING', 'QUEUED'],
-          taskIdentifier: screenshots_task_keys.take_metric_screenshot,
+      if (
+        await shouldTakeScreenshot({
           tag,
-          limit: 1,
+          key: screenshots_task_keys.take_metric_screenshot,
+          context: c,
         })
-        .then((res) => res.data[0]);
-
-      if (!lastTask && !is_screenshot) {
+      ) {
         const organizationId =
           (await getUserOrganizationId(user.id).then((res) => res?.organizationId)) || '';
         await tasks.trigger(
diff --git a/apps/server/src/shared-helpers/screenshots.ts b/apps/server/src/shared-helpers/screenshots.ts
new file mode 100644
index 000000000..9759e58f6
--- /dev/null
+++ b/apps/server/src/shared-helpers/screenshots.ts
@@ -0,0 +1,30 @@
+import type { screenshots_task_keys } from '@buster-app/trigger/task-keys';
+import { runs } from '@trigger.dev/sdk';
+import type { Context } from 'hono';
+
+export const shouldTakeScreenshot = async ({
+  tag,
+  key,
+  context,
+}: {
+  tag: string;
+  key: (typeof screenshots_task_keys)[keyof typeof screenshots_task_keys];
+  context: Context;
+}): Promise<boolean> => {
+  const hasReferrer = !!context.req.header('referer');
+
+  if (!hasReferrer) {
+    return false;
+  }
+
+  const lastTask = await runs
+    .list({
+      status: ['EXECUTING', 'QUEUED'],
+      taskIdentifier: key,
+      tag,
+      limit: 1,
+    })
+    .then((res) => res.data[0]);
+
+  return !lastTask;
+};
diff --git a/apps/web/src/api/buster_rest/metrics/requests.ts b/apps/web/src/api/buster_rest/metrics/requests.ts
index a8544bd26..cff87ad9b 100644
--- a/apps/web/src/api/buster_rest/metrics/requests.ts
+++ b/apps/web/src/api/buster_rest/metrics/requests.ts
@@ -5,12 +5,13 @@ import type {
   DeleteMetricResponse,
   DuplicateMetricRequest,
   DuplicateMetricResponse,
-  GetMetricDataRequest,
   GetMetricListRequest,
   GetMetricParams,
   GetMetricQuery,
   GetMetricResponse,
   ListMetricsResponse,
+  MetricDataParams,
+  MetricDataQuery,
   MetricDataResponse,
   MetricDownloadParams,
   MetricDownloadQueryParams,
@@ -40,7 +41,7 @@ export const getMetric = async ({
 export const getMetricData = async ({
   id,
   ...params
-}: GetMetricDataRequest & GetMetricParams): Promise<MetricDataResponse> => {
+}: MetricDataParams & MetricDataQuery): Promise<MetricDataResponse> => {
   return mainApiV2
     .get<MetricDataResponse>(`/metric_files/${id}/data`, { params })
     .then((res) => res.data);
diff --git a/packages/server-shared/src/metrics/requests.types.ts b/packages/server-shared/src/metrics/requests.types.ts
index f09fc1181..4d040c0c9 100644
--- a/packages/server-shared/src/metrics/requests.types.ts
+++ b/packages/server-shared/src/metrics/requests.types.ts
@@ -61,7 +61,6 @@ export const BulkUpdateMetricVerificationStatusRequestSchema = z.array(
 
 export type GetMetricParams = z.infer<typeof GetMetricParamsSchema>;
 export type GetMetricQuery = z.infer<typeof GetMetricQuerySchema>;
-export type GetMetricDataRequest = z.infer<typeof GetMetricDataRequestSchema>;
 export type GetMetricListRequest = z.infer<typeof GetMetricListRequestSchema>;
 export type UpdateMetricRequest = z.infer<typeof UpdateMetricRequestSchema>;
 export type DeleteMetricRequest = z.infer<typeof DeleteMetricRequestSchema>;
diff --git a/packages/server-shared/src/metrics/responses.types.ts b/packages/server-shared/src/metrics/responses.types.ts
index edffd40ac..7fa746b23 100644
--- a/packages/server-shared/src/metrics/responses.types.ts
+++ b/packages/server-shared/src/metrics/responses.types.ts
@@ -55,7 +55,6 @@ export const MetricDataQuerySchema = z.object({
   version_number: z.coerce.number().int().min(1).optional(),
   report_file_id: z.string().uuid().optional(),
   password: z.string().min(1).optional(),
-  is_screenshot: z.boolean().default(false).optional(),
 });
 
 export type MetricDataQuery = z.infer<typeof MetricDataQuerySchema>;