fix(capture): retry content-script message on one-click auto-capture

src/lib/capture.ts

@@ -53,6 +53,44 @@ async function ensureInjectable(tabId: number): Promise<void> {
   }
 }
 
+/**
+ * True for the transient error Chrome throws when a message is sent to a tab
+ * whose content script has not registered its listener yet.
+ */
+export function isNoReceiverError(error: unknown): boolean {
+  const message = error instanceof Error ? error.message : String(error);
+  return /receiving end does not exist|could not establish connection/i.test(message);
+}
+
+/**
+ * Send a message to a tab's content script, re-injecting it and retrying while
+ * the receiving end is not ready. This is the difference between manual capture
+ * (the user clicks seconds later, script is ready) and one-click auto-capture
+ * (fires immediately on editor load, before the listener has registered).
+ */
+export async function sendToContentScript<T>(
+  tabId: number,
+  message: unknown,
+  options: { retries?: number; delayMs?: number } = {}
+): Promise<T> {
+  const retries = options.retries ?? 6;
+  const delayMs = options.delayMs ?? 150;
+
+  let lastError: unknown;
+  for (let attempt = 0; attempt <= retries; attempt += 1) {
+    try {
+      return (await chrome.tabs.sendMessage(tabId, message)) as T;
+    } catch (error) {
+      lastError = error;
+      if (!isNoReceiverError(error)) throw error;
+      await ensureInjectable(tabId);
+      await wait(delayMs);
+    }
+  }
+
+  throw lastError instanceof Error ? lastError : new Error("Content script did not respond");
+}
+
 export async function captureFullPage(
   tabId: number,
   windowId: number,
@@ -67,7 +105,9 @@ export async function captureFullPage(
   try {
     await ensureInjectable(tabId);
 
-    const metrics = await sendMessage<PageMetrics>(tabId, { type: "SB_GET_PAGE_METRICS" });
+    // The first message races content-script startup on one-click auto-capture,
+    // so retry it (re-injecting) until the listener is ready.
+    const metrics = await sendToContentScript<PageMetrics>(tabId, { type: "SB_GET_PAGE_METRICS" });
     const steps = buildScrollSteps(metrics.fullHeight, metrics.viewportHeight);
 
     const segments: Array<{ y: number; dataUrl: string }> = [];

tests/capture.test.ts

@@ -1,5 +1,5 @@
-import { describe, expect, it } from "vitest";
-import { buildScrollSteps } from "../src/lib/capture";
+import { describe, expect, it, vi } from "vitest";
+import { buildScrollSteps, isNoReceiverError, sendToContentScript } from "../src/lib/capture";
 
 describe("buildScrollSteps", () => {
   it("returns single step when content fits viewport", () => {
@@ -14,3 +14,71 @@ describe("buildScrollSteps", () => {
     expect(buildScrollSteps(3000, 1000)).toEqual([0, 1000, 2000]);
   });
 });
+
+describe("isNoReceiverError", () => {
+  it("matches the content-script-not-ready errors", () => {
+    expect(
+      isNoReceiverError(new Error("Could not establish connection. Receiving end does not exist."))
+    ).toBe(true);
+    expect(isNoReceiverError(new Error("Receiving end does not exist"))).toBe(true);
+  });
+
+  it("does not match unrelated errors", () => {
+    expect(isNoReceiverError(new Error("Cannot access a chrome:// URL"))).toBe(false);
+  });
+});
+
+describe("sendToContentScript", () => {
+  function stubChrome(sendMessage: ReturnType<typeof vi.fn>) {
+    const executeScript = vi.fn().mockResolvedValue(undefined);
+    (globalThis as unknown as { chrome: unknown }).chrome = {
+      tabs: { sendMessage },
+      scripting: { executeScript }
+    };
+    return { executeScript };
+  }
+
+  it("returns the response on the first try", async () => {
+    const sendMessage = vi.fn().mockResolvedValue({ ok: true });
+    stubChrome(sendMessage);
+    await expect(sendToContentScript(1, { type: "X" }, { delayMs: 0 })).resolves.toEqual({
+      ok: true
+    });
+    expect(sendMessage).toHaveBeenCalledTimes(1);
+  });
+
+  it("re-injects and retries when the receiver is not ready, then succeeds", async () => {
+    const sendMessage = vi
+      .fn()
+      .mockRejectedValueOnce(
+        new Error("Could not establish connection. Receiving end does not exist.")
+      )
+      .mockRejectedValueOnce(new Error("Receiving end does not exist"))
+      .mockResolvedValue({ ok: true });
+    const { executeScript } = stubChrome(sendMessage);
+
+    await expect(sendToContentScript(1, { type: "X" }, { delayMs: 0 })).resolves.toEqual({
+      ok: true
+    });
+    expect(sendMessage).toHaveBeenCalledTimes(3);
+    expect(executeScript).toHaveBeenCalledTimes(2);
+  });
+
+  it("rethrows non-receiver errors immediately without retrying", async () => {
+    const sendMessage = vi.fn().mockRejectedValue(new Error("Cannot access a chrome:// URL"));
+    stubChrome(sendMessage);
+    await expect(sendToContentScript(1, { type: "X" }, { delayMs: 0 })).rejects.toThrow(
+      "chrome://"
+    );
+    expect(sendMessage).toHaveBeenCalledTimes(1);
+  });
+
+  it("gives up after exhausting retries", async () => {
+    const sendMessage = vi.fn().mockRejectedValue(new Error("Receiving end does not exist"));
+    stubChrome(sendMessage);
+    await expect(sendToContentScript(1, { type: "X" }, { retries: 2, delayMs: 0 })).rejects.toThrow(
+      "Receiving end does not exist"
+    );
+    expect(sendMessage).toHaveBeenCalledTimes(3);
+  });
+});