From 14c352de88f4c3200667afa4c0f124583142573a Mon Sep 17 00:00:00 2001 From: Jake Jarvis Date: Thu, 10 Sep 2026 12:02:40 -0400 Subject: [PATCH] fix: reuse screenshot workflow and improve client polling Prevent duplicate screenshot runs by introducing an ownership hook token and reusing active workflow runs. The workflow input now includes domainId and registers an ownership hook (getScreenshotWorkflowToken) to detect conflicts. API route: check for an active hook before starting, return no-store cache headers, and map workflow errors (including cancelled vs unavailable runs) and rate limits more clearly. Client: rewrite useScreenshot to a react-query based implementation with robust parsing, polling/backoff, reloadable image UI and placeholders. Added tests for the API and component behavior. --- apps/web/app/api/screenshot/route.test.ts | 109 ++++ apps/web/app/api/screenshot/route.ts | 76 ++- .../web/components/domain/screenshot.test.tsx | 152 +++++ apps/web/components/domain/screenshot.tsx | 527 ++++++++++-------- apps/web/workflows/screenshot/index.ts | 1 + apps/web/workflows/screenshot/workflow.ts | 17 +- 6 files changed, 638 insertions(+), 244 deletions(-) create mode 100644 apps/web/app/api/screenshot/route.test.ts create mode 100644 apps/web/components/domain/screenshot.test.tsx diff --git a/apps/web/app/api/screenshot/route.test.ts b/apps/web/app/api/screenshot/route.test.ts new file mode 100644 index 00000000..7996c042 --- /dev/null +++ b/apps/web/app/api/screenshot/route.test.ts @@ -0,0 +1,109 @@ +import { NextRequest } from "next/server"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + hookNotFound: new Error("hook not found"), + runNotFound: new Error("run not found"), + checkRateLimit: vi.fn<() => Promise>(), + getDomainById: vi.fn<(domainId: string) => Promise>(), + getHookByToken: vi.fn<(token: string) => Promise<{ runId: string }>>(), + getRun: vi.fn<(runId: string) => { status: Promise; returnValue?: Promise }>(), + getScreenshotByDomainId: vi.fn<(domainId: string) => Promise>(), + isDomainBlocked: vi.fn<(domain: string) => Promise>(), + start: vi.fn<(...args: unknown[]) => Promise<{ runId: string }>>(), +})); + +vi.mock("workflow/api", () => ({ + getHookByToken: mocks.getHookByToken, + getRun: mocks.getRun, + start: mocks.start, +})); +vi.mock("workflow/errors", () => ({ + HookNotFoundError: { is: (error: unknown) => error === mocks.hookNotFound }, + WorkflowRunNotFoundError: { is: (error: unknown) => error === mocks.runNotFound }, +})); +vi.mock("@/lib/ratelimit/api", () => ({ checkRateLimit: mocks.checkRateLimit })); +vi.mock("@domainstack/db/queries/blocked-domains", () => ({ + isDomainBlocked: mocks.isDomainBlocked, +})); +vi.mock("@domainstack/db/queries/domains", () => ({ getDomainById: mocks.getDomainById })); +vi.mock("@domainstack/db/queries/screenshots", () => ({ + getScreenshotByDomainId: mocks.getScreenshotByDomainId, +})); + +import { GET, POST } from "./route"; + +function postRequest() { + return new NextRequest("https://domainstack.io/api/screenshot", { + method: "POST", + body: JSON.stringify({ domainId: "domain-1" }), + headers: { "Content-Type": "application/json" }, + }); +} + +describe("screenshot API", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.checkRateLimit.mockResolvedValue({ success: true }); + mocks.getDomainById.mockResolvedValue({ id: "domain-1", name: "example.com" }); + mocks.getScreenshotByDomainId.mockResolvedValue(null); + mocks.getHookByToken.mockRejectedValue(mocks.hookNotFound); + mocks.start.mockResolvedValue({ runId: "run-new" }); + }); + + it("reuses the active workflow registered for a domain", async () => { + mocks.getHookByToken.mockResolvedValue({ runId: "run-active" }); + + const response = await POST(postRequest()); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual({ status: "running", runId: "run-active" }); + expect(mocks.getHookByToken).toHaveBeenCalledWith("screenshot:domain-1"); + expect(mocks.start).not.toHaveBeenCalled(); + expect(response.headers.get("Cache-Control")).toBe("no-cache, no-store"); + }); + + it("passes the domain identity into a newly started workflow", async () => { + const response = await POST(postRequest()); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual({ status: "running", runId: "run-new" }); + expect(mocks.start).toHaveBeenCalledWith(expect.any(Function), [ + { domain: "example.com", domainId: "domain-1" }, + ]); + }); + + it("reports a cancelled workflow as terminal", async () => { + mocks.getRun.mockReturnValue({ status: Promise.resolve("cancelled") }); + + const response = await GET( + new NextRequest("https://domainstack.io/api/screenshot?runId=run-cancelled"), + ); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual({ + status: "failed", + error: "workflow_cancelled", + }); + expect(response.headers.get("Cache-Control")).toBe("no-cache, no-store"); + }); + + it("distinguishes a not-yet-visible run from a status backend failure", async () => { + mocks.getRun.mockReturnValueOnce({ status: Promise.reject(mocks.runNotFound) }); + + const notFoundResponse = await GET( + new NextRequest("https://domainstack.io/api/screenshot?runId=run-pending"), + ); + expect(notFoundResponse.status).toBe(404); + await expect(notFoundResponse.json()).resolves.toEqual({ error: "Run not found" }); + + mocks.getRun.mockReturnValueOnce({ status: Promise.reject(new Error("world unavailable")) }); + const unavailableResponse = await GET( + new NextRequest("https://domainstack.io/api/screenshot?runId=run-pending"), + ); + expect(unavailableResponse.status).toBe(503); + await expect(unavailableResponse.json()).resolves.toEqual({ + error: "Workflow status unavailable", + }); + }); +}); diff --git a/apps/web/app/api/screenshot/route.ts b/apps/web/app/api/screenshot/route.ts index bd2c9cdb..a758bb20 100644 --- a/apps/web/app/api/screenshot/route.ts +++ b/apps/web/app/api/screenshot/route.ts @@ -1,9 +1,14 @@ import type { NextRequest } from "next/server"; import { NextResponse } from "next/server"; -import { getRun, start } from "workflow/api"; +import { getHookByToken, getRun, start } from "workflow/api"; +import { HookNotFoundError, WorkflowRunNotFoundError } from "workflow/errors"; import { checkRateLimit } from "@/lib/ratelimit/api"; -import { type ScreenshotWorkflowResult, screenshotWorkflow } from "@/workflows/screenshot"; +import { + getScreenshotWorkflowToken, + type ScreenshotWorkflowResult, + screenshotWorkflow, +} from "@/workflows/screenshot"; import { isDomainBlocked } from "@domainstack/db/queries/blocked-domains"; import { getDomainById } from "@domainstack/db/queries/domains"; import { getScreenshotByDomainId } from "@domainstack/db/queries/screenshots"; @@ -22,6 +27,16 @@ type ScreenshotStatusResponse = | { status: "completed"; cached: false; success: false; error: string; data: { url: null } } | { status: "failed"; error: string }; +const NO_STORE_HEADERS = { + "Cache-Control": "no-cache, no-store", +} as const; + +function withNoStore(headers?: HeadersInit): Headers { + const result = new Headers(headers); + result.set("Cache-Control", NO_STORE_HEADERS["Cache-Control"]); + return result; +} + /** * POST /api/screenshot * @@ -101,13 +116,35 @@ export async function POST( cached: true, data: { url: cachedScreenshot.url, blocked }, }, - { headers: rateLimit.headers }, + { headers: withNoStore(rateLimit.headers) }, ); } } - // Cache miss - start workflow - const run = await start(screenshotWorkflow, [{ domain: domain.name }]); + // Cache miss - reuse an active workflow when its ownership hook has + // already been registered. The workflow also checks for hook conflicts + // to close the race between this advisory lookup and start(). + const token = getScreenshotWorkflowToken(domainId); + try { + const activeHook = await getHookByToken(token); + logger.debug( + { domainId, domain: domain.name, runId: activeHook.runId }, + "reusing active screenshot workflow", + ); + return NextResponse.json( + { status: "running", runId: activeHook.runId }, + { headers: withNoStore(rateLimit.headers) }, + ); + } catch (err) { + if (!HookNotFoundError.is(err)) { + logger.warn( + { err, domainId, domain: domain.name }, + "failed to look up screenshot workflow", + ); + } + } + + const run = await start(screenshotWorkflow, [{ domain: domain.name, domainId }]); logger.debug( { domainId, domain: domain.name, runId: run.runId }, @@ -119,7 +156,7 @@ export async function POST( status: "running", runId: run.runId, }, - { headers: rateLimit.headers }, + { headers: withNoStore(rateLimit.headers) }, ); } catch (err) { logger.error({ err }, "failed to start screenshot workflow"); @@ -172,24 +209,35 @@ export async function GET( data: result.data, ...(!result.success && { error: result.error }), } as ScreenshotStatusResponse, - { headers: rateLimit.headers }, + { headers: withNoStore(rateLimit.headers) }, ); } - if (status === "failed") { + if (status === "failed" || status === "cancelled") { return NextResponse.json( { status: "failed", - error: "workflow_failed", + error: status === "cancelled" ? "workflow_cancelled" : "workflow_failed", }, - { headers: rateLimit.headers }, + { headers: withNoStore(rateLimit.headers) }, ); } - // Still running - return NextResponse.json({ status: "running" }, { headers: rateLimit.headers }); + // Pending and running are both non-terminal. + return NextResponse.json({ status: "running" }, { headers: withNoStore(rateLimit.headers) }); } catch (err) { - logger.debug({ err, runId }, "workflow run unavailable"); - return NextResponse.json({ error: "Run not found" }, { status: 404 }); + if (WorkflowRunNotFoundError.is(err)) { + logger.debug({ err, runId }, "workflow run not visible yet"); + return NextResponse.json( + { error: "Run not found" }, + { status: 404, headers: NO_STORE_HEADERS }, + ); + } + + logger.warn({ err, runId }, "failed to get workflow run status"); + return NextResponse.json( + { error: "Workflow status unavailable" }, + { status: 503, headers: NO_STORE_HEADERS }, + ); } } diff --git a/apps/web/components/domain/screenshot.test.tsx b/apps/web/components/domain/screenshot.test.tsx new file mode 100644 index 00000000..23aa2699 --- /dev/null +++ b/apps/web/components/domain/screenshot.test.tsx @@ -0,0 +1,152 @@ +import { QueryClientProvider } from "@tanstack/react-query"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { page } from "vitest/browser"; + +vi.mock("@/lib/analytics/client", () => ({ + analytics: { + track: vi.fn<(event: string, properties?: Record) => void>(), + trackException: vi.fn<(error: unknown, context?: Record) => void>(), + }, +})); +vi.mock("sonner", () => ({ + toast: { error: vi.fn<(message: string, options?: unknown) => void>() }, +})); + +import { createTestQueryClient, render, renderHook } from "@/mocks/react"; + +import { Screenshot, useScreenshot } from "./screenshot"; + +const screenshotUrl = "https://example.public.blob.vercel-storage.com/screenshot.webp"; + +function jsonResponse(body: unknown, init?: ResponseInit) { + const headers = new Headers(init?.headers); + headers.set("Content-Type", "application/json"); + return new Response(JSON.stringify(body), { + ...init, + headers, + }); +} + +describe("useScreenshot", () => { + beforeEach(() => { + vi.stubGlobal("fetch", vi.fn()); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("recovers after an early status lookup fails", async () => { + const fetchMock = vi.mocked(fetch); + fetchMock + .mockResolvedValueOnce(jsonResponse({ status: "running", runId: "run-1" })) + .mockResolvedValueOnce(jsonResponse({ error: "Run not found" }, { status: 404 })) + .mockResolvedValueOnce( + jsonResponse({ + status: "completed", + cached: false, + success: true, + data: { url: screenshotUrl }, + }), + ); + + const queryClient = createTestQueryClient(); + const view = await renderHook( + () => useScreenshot({ domain: "example.com", domainId: "domain-1" }), + { + wrapper: ({ children }) => ( + {children} + ), + }, + ); + + await vi.waitFor(() => + expect(queryClient.getQueryData(["screenshot", "domain-1"])).toEqual({ + status: "running", + runId: "run-1", + }), + ); + await queryClient.refetchQueries({ queryKey: ["screenshot", "domain-1"] }); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + expect(view.result.current.isLoading).toBe(true); + + await queryClient.refetchQueries({ queryKey: ["screenshot", "domain-1"] }); + await vi.waitFor(() => + expect(queryClient.getQueryData(["screenshot", "domain-1"])).toMatchObject({ + status: "completed", + data: { url: screenshotUrl }, + }), + ); + await vi.waitFor(() => expect(view.result.current.data?.url).toBe(screenshotUrl)); + expect(view.result.current.isLoading).toBe(false); + }); + + it("shares an active run across observers and remounts", async () => { + const fetchMock = vi.mocked(fetch); + fetchMock + .mockResolvedValueOnce(jsonResponse({ status: "running", runId: "run-shared" })) + .mockResolvedValue( + jsonResponse({ + status: "completed", + cached: false, + success: true, + data: { url: screenshotUrl }, + }), + ); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: { children: React.ReactNode }) => ( + {children} + ); + const first = await renderHook( + () => useScreenshot({ domain: "example.com", domainId: "domain-1" }), + { wrapper }, + ); + await vi.waitFor(() => + expect(queryClient.getQueryData(["screenshot", "domain-1"])).toEqual({ + status: "running", + runId: "run-shared", + }), + ); + + await first.unmount(); + const second = await renderHook( + () => useScreenshot({ domain: "example.com", domainId: "domain-1" }), + { wrapper }, + ); + + await vi.waitFor(() => + expect(queryClient.getQueryData(["screenshot", "domain-1"])).toMatchObject({ + status: "completed", + data: { url: screenshotUrl }, + }), + ); + await vi.waitFor(() => expect(second.result.current.data?.url).toBe(screenshotUrl)); + const postRequests = fetchMock.mock.calls.filter(([, init]) => init?.method === "POST"); + expect(postRequests).toHaveLength(1); + }); +}); + +describe("Screenshot", () => { + it("lets the user reload an image that failed without refreshing the page", async () => { + await render( + , + ); + + const image = page.getByRole("img", { name: "Homepage preview of example.com" }); + await expect.element(image).toBeInTheDocument(); + image.element().dispatchEvent(new Event("error")); + const reloadButton = page.getByRole("button", { name: "Reload preview" }); + await expect.element(reloadButton).toBeInTheDocument(); + + await reloadButton.click(); + const reloadedImage = page + .getByRole("img", { name: "Homepage preview of example.com" }) + .element() as HTMLImageElement; + expect(reloadedImage.src).toContain("domainstack-reload=1"); + }); +}); diff --git a/apps/web/components/domain/screenshot.tsx b/apps/web/components/domain/screenshot.tsx index 6a569507..e1fed72c 100644 --- a/apps/web/components/domain/screenshot.tsx +++ b/apps/web/components/domain/screenshot.tsx @@ -1,9 +1,9 @@ "use client"; import { IconCircleX, IconShieldExclamation } from "@tabler/icons-react"; -import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; +import { useQuery, useQueryClient } from "@tanstack/react-query"; import Image from "next/image"; -import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { useEffect, useRef, useState } from "react"; import { toast } from "sonner"; import { analytics } from "@/lib/analytics/client"; @@ -13,69 +13,81 @@ import { Spinner } from "@domainstack/ui/spinner"; import { cn } from "@domainstack/ui/utils"; const POLL_INTERVAL_MS = 2000; +const POLL_RECOVERY_INTERVAL_MS = 5000; -type StartParseResult = - | { status: "completed"; cached: true; data: ScreenshotData } +type ScreenshotQueryState = + | { status: "completed"; source: "cache" | "workflow"; data: ScreenshotData } | { status: "running"; runId: string } - | { status: "error"; error: string } - | { status: "rate_limited"; retryAfter: number }; - -type StatusParseResult = - | { status: "running" } - | { status: "completed"; data: ScreenshotData } | { status: "failed"; error: string } - | { status: "error"; error: string } - | { status: "rate_limited"; retryAfter: number }; + | { status: "rate_limited"; retryAfter: number; runId?: string }; -function parseStartResponse(raw: unknown): StartParseResult { +type TerminalScreenshotQueryState = Extract< + ScreenshotQueryState, + { status: "completed" | "failed" } +>; + +function getScreenshotQueryKey(domain: string, domainId?: string) { + return ["screenshot", domainId ?? domain] as const; +} + +function parseScreenshotData(raw: unknown): ScreenshotData { if (!raw || typeof raw !== "object") { - return { status: "error", error: "Invalid response" }; + throw new Error("Screenshot response is missing data"); + } + + const data = raw as Record; + return { + url: typeof data.url === "string" ? data.url : null, + blocked: data.blocked === true, + }; +} + +function parseStartResponse(raw: unknown): ScreenshotQueryState { + if (!raw || typeof raw !== "object") { + throw new Error("Invalid screenshot response"); } const obj = raw as Record; if ("error" in obj && !("status" in obj)) { - return { - status: "error", - error: typeof obj.error === "string" ? obj.error : "Invalid response", - }; + throw new Error(typeof obj.error === "string" ? obj.error : "Screenshot request failed"); } if (obj.status === "running" && typeof obj.runId === "string") { return { status: "running", runId: obj.runId }; } - if (obj.status === "completed" && obj.data) { - const data = obj.data as Record; + if (obj.status === "completed" && obj.success === false) { return { - status: "completed", - cached: true, - data: { - url: typeof data.url === "string" ? data.url : null, - blocked: data.blocked === true, - }, + status: "failed", + error: typeof obj.error === "string" ? obj.error : "Screenshot capture failed", }; } - return { status: "error", error: "Unknown response format" }; + if (obj.status === "completed" && obj.data) { + return { + status: "completed", + source: "cache", + data: parseScreenshotData(obj.data), + }; + } + + throw new Error("Unknown screenshot response format"); } -function parseStatusResponse(raw: unknown): StatusParseResult { +function parseStatusResponse(raw: unknown, runId: string): ScreenshotQueryState { if (!raw || typeof raw !== "object") { - return { status: "error", error: "Invalid response" }; + throw new Error("Invalid screenshot status response"); } const obj = raw as Record; if ("error" in obj && !("status" in obj)) { - return { - status: "error", - error: typeof obj.error === "string" ? obj.error : "Invalid response", - }; + throw new Error(typeof obj.error === "string" ? obj.error : "Screenshot status unavailable"); } if (obj.status === "running") { - return { status: "running" }; + return { status: "running", runId }; } if (obj.status === "failed") { @@ -86,17 +98,73 @@ function parseStatusResponse(raw: unknown): StatusParseResult { } if (obj.status === "completed" && obj.data) { - const data = obj.data as Record; return { status: "completed", - data: { - url: typeof data.url === "string" ? data.url : null, - blocked: data.blocked === true, - }, + source: "workflow", + data: parseScreenshotData(obj.data), }; } - return { status: "error", error: "Unknown response format" }; + throw new Error("Unknown screenshot status response format"); +} + +async function readErrorMessage(response: Response, fallback: string): Promise { + try { + const raw = (await response.json()) as { error?: unknown }; + return typeof raw.error === "string" ? raw.error : fallback; + } catch { + return fallback; + } +} + +async function startScreenshot(domainId: string): Promise { + const response = await fetch("/api/screenshot", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ domainId }), + }); + + if (response.status === 429) { + return { status: "rate_limited", retryAfter: parseRetryAfterHeader(response) }; + } + + if (!response.ok) { + const error = await readErrorMessage(response, `Screenshot request failed: ${response.status}`); + if (response.status >= 400 && response.status < 500) { + return { status: "failed", error }; + } + throw new Error(error); + } + + return parseStartResponse(await response.json()); +} + +async function pollScreenshot(runId: string): Promise { + const response = await fetch(`/api/screenshot?runId=${encodeURIComponent(runId)}`, { + cache: "no-store", + }); + + if (response.status === 429) { + return { + status: "rate_limited", + retryAfter: parseRetryAfterHeader(response), + runId, + }; + } + + if (!response.ok) { + throw new Error( + await readErrorMessage(response, `Screenshot status poll failed: ${response.status}`), + ); + } + + return parseStatusResponse(await response.json(), runId); +} + +function isTerminalState( + state: ScreenshotQueryState | undefined, +): state is TerminalScreenshotQueryState { + return state?.status === "completed" || state?.status === "failed"; } export interface UseScreenshotResult { @@ -106,6 +174,118 @@ export interface UseScreenshotResult { hasFailed: boolean; } +function ScreenshotPlaceholder({ + isLoading, + blocked, + aspectClassName, + onReload, +}: { + isLoading: boolean; + blocked: boolean; + aspectClassName: string; + onReload?: () => void; +}) { + return ( +
+
+ {isLoading ? ( + <> + + Taking screenshot… + + ) : blocked ? ( + <> + + Screenshot unavailable for this domain. + + ) : ( + <> + + Unable to take a screenshot. + {onReload ? ( + + ) : null} + + )} +
+
+ ); +} + +function withReloadToken(url: string, reloadCount: number): string { + if (reloadCount === 0) return url; + + try { + const reloadUrl = new URL(url); + reloadUrl.searchParams.set("domainstack-reload", String(reloadCount)); + return reloadUrl.toString(); + } catch { + const separator = url.includes("?") ? "&" : "?"; + return `${url}${separator}domainstack-reload=${reloadCount}`; + } +} + +function ScreenshotImage({ + domain, + url, + width, + height, + imageClassName, + aspectClassName, +}: { + domain: string; + url: string; + width: number; + height: number; + imageClassName?: string; + aspectClassName: string; +}) { + const [reloadCount, setReloadCount] = useState(0); + const [hasFailed, setHasFailed] = useState(false); + + if (hasFailed) { + return ( + { + setReloadCount((current) => current + 1); + setHasFailed(false); + }} + /> + ); + } + + return ( + + {`Homepage setHasFailed(true)} + /> + + ); +} + /** * Hook to fetch a screenshot for a domain. * Call this in a component that stays mounted to keep polling active. @@ -120,174 +300,92 @@ export function useScreenshot({ enabled?: boolean; }): UseScreenshotResult { const queryClient = useQueryClient(); - const [runId, setRunId] = useState(null); - const [screenshotData, setScreenshotData] = useState(null); - const hasStartedRef = useRef(false); - const startedForDomainRef = useRef(null); - const [rateLimitedUntil, setRateLimitedUntil] = useState(null); - const retryTimeoutRef = useRef | null>(null); + const queryKey = getScreenshotQueryKey(domain, domainId); + const screenshotQuery = useQuery({ + queryKey, + queryFn: async () => { + const current = queryClient.getQueryData(queryKey); - const screenshotQueryKey = useMemo(() => ["screenshot", domain], [domain]); - const cachedData = queryClient.getQueryData(screenshotQueryKey); - - const [trackedDomain, setTrackedDomain] = useState(domain); - if (domain !== trackedDomain) { - setTrackedDomain(domain); - setScreenshotData(null); - setRunId(null); - setRateLimitedUntil(null); - } - - const startScreenshot = useCallback(async (id: string) => { - const response = await fetch("/api/screenshot", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ domainId: id }), - }); - - if (response.status === 429) { - const retryAfter = parseRetryAfterHeader(response); - return { status: "rate_limited", retryAfter } as const; - } - - if (!response.ok) { - throw new Error(`Screenshot request failed: ${response.status}`); - } - - return parseStartResponse(await response.json()); - }, []); - - const startMutation = useMutation({ - mutationFn: startScreenshot, - onSuccess: (data) => { - if (data.status === "completed") { - setScreenshotData(data.data); - queryClient.setQueryData(screenshotQueryKey, data.data); - analytics.track("screenshot_loaded_from_cache", { domain }); - } else if (data.status === "running") { - setRunId(data.runId); - analytics.track("screenshot_requested", { domain }); - } else if (data.status === "rate_limited") { - const retryAt = Date.now() + data.retryAfter * 1000; - setRateLimitedUntil(retryAt); - toast.error("Too many requests", { - description: `Please wait ${data.retryAfter} second${data.retryAfter !== 1 ? "s" : ""} before trying again.`, - }); - analytics.track("screenshot_rate_limited", { - domain, - retryAfter: data.retryAfter, - }); - - if (retryTimeoutRef.current) { - clearTimeout(retryTimeoutRef.current); - } - retryTimeoutRef.current = setTimeout(() => { - hasStartedRef.current = false; - setRateLimitedUntil(null); - }, data.retryAfter * 1000); + if (current?.status === "running") { + return pollScreenshot(current.runId); } - }, - onError: (error) => { - analytics.trackException(error, { domain }); - }, - }); - - const statusQuery = useQuery({ - queryKey: ["screenshot-status", runId], - queryFn: async (): Promise => { - const response = await fetch(`/api/screenshot?runId=${runId}`); - - if (response.status === 429) { - const retryAfter = parseRetryAfterHeader(response); - return { status: "rate_limited", retryAfter }; + if (current?.status === "rate_limited" && current.runId) { + return pollScreenshot(current.runId); + } + if (isTerminalState(current)) { + return current; } - if (!response.ok) { - throw new Error(`Screenshot status poll failed: ${response.status}`); + if (!domainId) { + throw new Error("Screenshot domain ID is missing"); } - return parseStatusResponse(await response.json()); + return startScreenshot(domainId); }, - enabled: !!runId, - staleTime: Infinity, + enabled: enabled && !!domainId, + retry: false, + staleTime: (query) => (isTerminalState(query.state.data) ? Number.POSITIVE_INFINITY : 0), + refetchOnMount: (query) => !isTerminalState(query.state.data), refetchInterval: (query) => { - const { data } = query.state; - if (data?.status !== "running") { - if (data?.status === "rate_limited") { - return data.retryAfter * 1000; - } + const state = query.state.data; + if (isTerminalState(state)) { return false; } - return POLL_INTERVAL_MS; + if (state?.status === "rate_limited") { + return state.retryAfter * 1000; + } + return state?.status === "running" ? POLL_INTERVAL_MS : POLL_RECOVERY_INTERVAL_MS; }, + refetchIntervalInBackground: true, }); - // Side effects only: cache the completed screenshot and notify. Terminal - // poll status is derived from the query during render so we don't copy it - // into state or clear runId (that would change the query key and drop data). + const reportedStateRef = useRef(null); useEffect(() => { - const data = statusQuery.data; - if (!data || data.status === "running") return; + const state = screenshotQuery.data; + if (!state) return; - if (data.status === "rate_limited") { + const marker = + state.status === "running" + ? `running:${state.runId}` + : state.status === "completed" + ? `completed:${state.source}:${state.data.url ?? "none"}` + : state.status === "rate_limited" + ? `rate-limited:${state.runId ?? "start"}:${state.retryAfter}` + : `failed:${state.error}`; + if (reportedStateRef.current === marker) return; + reportedStateRef.current = marker; + + if (state.status === "running") { + analytics.track("screenshot_requested", { domain }); + } else if (state.status === "completed") { + analytics.track( + state.source === "cache" ? "screenshot_loaded_from_cache" : "screenshot_loaded_from_api", + { domain }, + ); + } else if (state.status === "rate_limited") { toast.error("Too many requests", { - description: `Polling paused. Retrying in ${data.retryAfter} seconds.`, + id: `screenshot-rate-limited-${domainId ?? domain}`, + description: `Retrying in ${state.retryAfter} second${state.retryAfter !== 1 ? "s" : ""}.`, + }); + analytics.track("screenshot_rate_limited", { + domain, + retryAfter: state.retryAfter, }); - return; } + }, [screenshotQuery.data, domain, domainId]); - if (data.status === "completed") { - queryClient.setQueryData(screenshotQueryKey, data.data); - analytics.track("screenshot_loaded_from_api", { domain }); - } - }, [statusQuery.data, queryClient, screenshotQueryKey, domain]); - - // Cleanup retry timeout on unmount - useEffect( - () => () => { - if (retryTimeoutRef.current) { - clearTimeout(retryTimeoutRef.current); - } - }, - [], - ); - - // Reset and auto-start when domain/enabled/domainId changes useEffect(() => { - if (startedForDomainRef.current !== domain) { - hasStartedRef.current = false; - startedForDomainRef.current = domain; - if (retryTimeoutRef.current) { - clearTimeout(retryTimeoutRef.current); - retryTimeoutRef.current = null; - } + if (screenshotQuery.error) { + analytics.trackException(screenshotQuery.error, { domain }); } + }, [screenshotQuery.error, domain]); - if (hasStartedRef.current || !enabled || !domainId || cachedData || screenshotData) { - return; - } + const state = screenshotQuery.data; + const data = state?.status === "completed" ? state.data : null; + const hasFailed = state?.status === "failed"; + const error = hasFailed ? new Error(state.error) : (screenshotQuery.error ?? null); + const isLoading = enabled && (!domainId || (!data && !hasFailed)); - if (rateLimitedUntil && Date.now() < rateLimitedUntil) { - return; - } - - hasStartedRef.current = true; - startMutation.mutate(domainId); - }, [domain, enabled, domainId, cachedData, screenshotData, startMutation, rateLimitedUntil]); - - // Derive return values - const polledData = statusQuery.data?.status === "completed" ? statusQuery.data.data : undefined; - const finalData = screenshotData ?? polledData ?? cachedData ?? null; - const pollError = statusQuery.data?.status === "error" ? new Error(statusQuery.data.error) : null; - const error = startMutation.error ?? statusQuery.error ?? pollError; - const hasFailed = statusQuery.data?.status === "failed"; - const isLoading = - !finalData && - !error && - !hasFailed && - enabled && - (domainId === undefined || startMutation.isPending || !!runId); - - return { data: finalData, isLoading, error, hasFailed }; + return { data, isLoading, error, hasFailed }; } /** @@ -313,54 +411,27 @@ export function Screenshot({ imageClassName?: string; aspectClassName?: string; }) { - const [failedUrl, setFailedUrl] = useState(null); - const url = data?.url ?? null; const blocked = data?.blocked ?? false; return (
- {url && failedUrl !== url ? ( - - {`Homepage setFailedUrl(url)} - /> - + {url ? ( + ) : ( -
-
- {isLoading ? ( - <> - - Taking screenshot… - - ) : blocked ? ( - <> - - Screenshot unavailable for this domain. - - ) : ( - <> - - Unable to take a screenshot. - - )} -
-
+ )}
); diff --git a/apps/web/workflows/screenshot/index.ts b/apps/web/workflows/screenshot/index.ts index ee69bf7a..ef133121 100644 --- a/apps/web/workflows/screenshot/index.ts +++ b/apps/web/workflows/screenshot/index.ts @@ -1,4 +1,5 @@ export { + getScreenshotWorkflowToken, type ScreenshotWorkflowInput, type ScreenshotWorkflowResult, screenshotWorkflow, diff --git a/apps/web/workflows/screenshot/workflow.ts b/apps/web/workflows/screenshot/workflow.ts index 094887cc..2f0e2cbc 100644 --- a/apps/web/workflows/screenshot/workflow.ts +++ b/apps/web/workflows/screenshot/workflow.ts @@ -1,12 +1,17 @@ -import { FatalError, RetryableError } from "workflow"; +import { createHook, FatalError, RetryableError } from "workflow"; import { checkBlocklist } from "@/workflows/shared/check-blocklist"; const VIEWPORT_WIDTH = 1200; const VIEWPORT_HEIGHT = 630; +export function getScreenshotWorkflowToken(domainId: string): string { + return `screenshot:${domainId}`; +} + export interface ScreenshotWorkflowInput { domain: string; + domainId: string; } export interface ScreenshotWorkflowData { @@ -50,7 +55,15 @@ export async function screenshotWorkflow( ): Promise { "use workflow"; - const { domain } = input; + const { domain, domainId } = input; + + using ownership = createHook({ + token: getScreenshotWorkflowToken(domainId), + }); + const conflictingRun = await ownership.getConflict(); + if (conflictingRun) { + return (await conflictingRun.returnValue) as ScreenshotWorkflowResult; + } // Step 1: Check if domain is blocked (shared step) const isBlocked = await checkBlocklist(domain);