diff --git a/src/lib/capture.ts b/src/lib/capture.ts index ff17121..3ac853e 100644 --- a/src/lib/capture.ts +++ b/src/lib/capture.ts @@ -53,6 +53,44 @@ async function ensureInjectable(tabId: number): Promise { } } +/** + * 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( + tabId: number, + message: unknown, + options: { retries?: number; delayMs?: number } = {} +): Promise { + 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(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(tabId, { type: "SB_GET_PAGE_METRICS" }); const steps = buildScrollSteps(metrics.fullHeight, metrics.viewportHeight); const segments: Array<{ y: number; dataUrl: string }> = []; diff --git a/tests/capture.test.ts b/tests/capture.test.ts index d25fb94..d900446 100644 --- a/tests/capture.test.ts +++ b/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) { + 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); + }); +});