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

@@ -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<unknown>>(),
+  getDomainById: vi.fn<(domainId: string) => Promise<unknown>>(),
+  getHookByToken: vi.fn<(token: string) => Promise<{ runId: string }>>(),
+  getRun: vi.fn<(runId: string) => { status: Promise<string>; returnValue?: Promise<unknown> }>(),
+  getScreenshotByDomainId: vi.fn<(domainId: string) => Promise<unknown>>(),
+  isDomainBlocked: vi.fn<(domain: string) => Promise<boolean>>(),
+  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",
+    });
+  });
+});

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 },
+    );
   }
 }

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<string, unknown>) => void>(),
+    trackException: vi.fn<(error: unknown, context?: Record<string, unknown>) => 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 }) => (
+          <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+        ),
+      },
+    );
+
+    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 }) => (
+      <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+    );
+    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(
+      <Screenshot
+        domain="example.com"
+        data={{ url: screenshotUrl, blocked: false }}
+        isLoading={false}
+      />,
+    );
+
+    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");
+  });
+});

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<string, unknown>;
+  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<string, unknown>;
 
   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<string, unknown>;
+  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<string, unknown>;
 
   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<string, unknown>;
     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<string> {
+  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<ScreenshotQueryState> {
+  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<ScreenshotQueryState> {
+  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 (
+    <div
+      className={`h-auto w-full ${aspectClassName} flex items-center justify-center bg-muted/50`}
+    >
+      <div
+        className="flex items-center gap-2 text-xs text-muted-foreground [&_svg]:size-4"
+        aria-live="polite"
+      >
+        {isLoading ? (
+          <>
+            <Spinner />
+            Taking screenshot…
+          </>
+        ) : blocked ? (
+          <>
+            <IconShieldExclamation />
+            Screenshot unavailable for this domain.
+          </>
+        ) : (
+          <>
+            <IconCircleX />
+            <span>Unable to take a screenshot.</span>
+            {onReload ? (
+              <button
+                type="button"
+                className="min-h-6 rounded-sm px-1.5 font-medium text-foreground underline underline-offset-2 hover:text-foreground/80 focus-visible:outline-2 focus-visible:outline-offset-2"
+                onClick={onReload}
+              >
+                Reload preview
+              </button>
+            ) : null}
+          </>
+        )}
+      </div>
+    </div>
+  );
+}
+
+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 (
+      <ScreenshotPlaceholder
+        isLoading={false}
+        blocked={false}
+        aspectClassName={aspectClassName}
+        onReload={() => {
+          setReloadCount((current) => current + 1);
+          setHasFailed(false);
+        }}
+      />
+    );
+  }
+
+  return (
+    <a href={`https://${domain}`} target="_blank" rel="noopener">
+      <Image
+        key={reloadCount}
+        src={withReloadToken(url, reloadCount)}
+        alt={`Homepage preview of ${domain}`}
+        width={width}
+        height={height}
+        className={cn("h-auto w-full object-cover", aspectClassName, imageClassName)}
+        unoptimized
+        priority={false}
+        draggable={false}
+        onError={() => setHasFailed(true)}
+      />
+    </a>
+  );
+}
+
 /**
  * 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<string | null>(null);
-  const [screenshotData, setScreenshotData] = useState<ScreenshotData | null>(null);
-  const hasStartedRef = useRef(false);
-  const startedForDomainRef = useRef<string | null>(null);
-  const [rateLimitedUntil, setRateLimitedUntil] = useState<number | null>(null);
-  const retryTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const queryKey = getScreenshotQueryKey(domain, domainId);
+  const screenshotQuery = useQuery<ScreenshotQueryState>({
+    queryKey,
+    queryFn: async () => {
+      const current = queryClient.getQueryData<ScreenshotQueryState>(queryKey);
 
-  const screenshotQueryKey = useMemo(() => ["screenshot", domain], [domain]);
-  const cachedData = queryClient.getQueryData<ScreenshotData>(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<StatusParseResult> => {
-      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<string | null>(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<string | null>(null);
-
   const url = data?.url ?? null;
   const blocked = data?.blocked ?? false;
 
   return (
     <div className={className}>
-      {url && failedUrl !== url ? (
-        <a href={`https://${domain}`} target="_blank" rel="noopener">
-          <Image
-            key={url}
-            src={url}
-            alt={`Homepage preview of ${domain}`}
-            width={width}
-            height={height}
-            className={cn("h-auto w-full object-cover", aspectClassName, imageClassName)}
-            unoptimized
-            priority={false}
-            draggable={false}
-            onError={() => setFailedUrl(url)}
-          />
-        </a>
+      {url ? (
+        <ScreenshotImage
+          key={url}
+          domain={domain}
+          url={url}
+          width={width}
+          height={height}
+          imageClassName={imageClassName}
+          aspectClassName={aspectClassName}
+        />
       ) : (
-        <div
-          className={`h-auto w-full ${aspectClassName} flex items-center justify-center bg-muted/50`}
-        >
-          <div
-            className="flex items-center gap-2 text-xs text-muted-foreground [&_svg]:size-4"
-            aria-live="polite"
-          >
-            {isLoading ? (
-              <>
-                <Spinner />
-                Taking screenshot…
-              </>
-            ) : blocked ? (
-              <>
-                <IconShieldExclamation />
-                Screenshot unavailable for this domain.
-              </>
-            ) : (
-              <>
-                <IconCircleX />
-                Unable to take a screenshot.
-              </>
-            )}
-          </div>
-        </div>
+        <ScreenshotPlaceholder
+          isLoading={isLoading}
+          blocked={blocked}
+          aspectClassName={aspectClassName}
+        />
       )}
     </div>
   );

apps/web/workflows/screenshot/index.ts

@@ -1,4 +1,5 @@
 export {
+  getScreenshotWorkflowToken,
   type ScreenshotWorkflowInput,
   type ScreenshotWorkflowResult,
   screenshotWorkflow,

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<ScreenshotWorkflowResult> {
   "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);