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