diff --git a/apps/extension/src/tools/vom/__tests__/capture-coordinator.test.ts b/apps/extension/src/tools/vom/__tests__/capture-coordinator.test.ts index 0bff254..624a2ed 100644 --- a/apps/extension/src/tools/vom/__tests__/capture-coordinator.test.ts +++ b/apps/extension/src/tools/vom/__tests__/capture-coordinator.test.ts @@ -170,6 +170,104 @@ describe("captureObservationFacts", () => { ).toBe(false); }); + it.each([ + "cancel", + "worker-abort", + ] as const)("joins target cleanup and stops claiming targets after %s", async (mode) => { + const frames = Array.from({ length: 5 }, (_, i) => ({ + frameId: `f${i}`, + target: { tabId: 4, ...(i ? { sessionId: `s${i}` } : {}) }, + })); + const { cdp } = fixture({ frames }); + const original = cdp.sendToTarget!; + const pending = new Map void; reject: (error: Error) => void }>(); + const calls: Array<{ target: string; method: string }> = []; + let release: (() => void) | undefined; + let released = false; + const send: NonNullable = async ( + target: CdpTarget, + method: string, + params?: object, + ) => { + calls.push({ target: target.sessionId ?? "main", method }); + if (method === "DOMSnapshot.captureSnapshot" && target.sessionId) { + await new Promise((resolve, reject) => + pending.set(target.sessionId!, { resolve, reject }), + ); + } + if (method === "Runtime.releaseObjectGroup") { + await new Promise((resolve) => { + release = resolve; + }); + released = true; + } + const result = await original(target, method, params); + if (method === "DOMSnapshot.captureSnapshot" && !target.sessionId) + (result as { strings: string[] }).strings[2] = "input"; + return result; + }; + cdp.sendToTarget = send; + cdp.send = (tabId, method, params) => send({ tabId }, method, params); + const controller = new AbortController(); + let settled = false; + const capture = captureObservationFacts(cdp, 4, controller.signal); + const rejected = expect(capture).rejects.toMatchObject({ name: "AbortError" }); + void capture.then( + () => { + settled = true; + }, + () => { + settled = true; + }, + ); + await vi.waitFor(() => { + expect(pending.size).toBe(3); + expect(release).toBeDefined(); + }); + if (mode === "cancel") controller.abort(); + // Cancellation may race with an ordinary transport rejection. Without an + // external signal, an escaping AbortError still terminates the worker pool. + pending + .get("s1")! + .reject( + mode === "cancel" + ? new Error("transport closed") + : new DOMException("target aborted", "AbortError"), + ); + pending.get("s2")!.resolve(); + pending.get("s3")!.resolve(); + await vi.waitFor(() => { + if (mode === "worker-abort") + expect( + calls.some( + (call) => call.target === "s3" && call.method === "Accessibility.getFullAXTree", + ), + ).toBe(true); + else + expect( + calls.some( + (call) => call.target === "s3" && call.method === "DOMSnapshot.captureSnapshot", + ), + ).toBe(true); + }); + // Allow worker rejection and the outer promise chain to run while cleanup + // is deliberately held. The old fail-fast pool settles here. + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(settled).toBe(false); + expect(released).toBe(false); + expect(calls.some((call) => call.target === "s4")).toBe(false); + if (mode === "cancel") + expect( + calls.some( + (call) => call.method === "Accessibility.enable" || call.method === "DOM.getDocument", + ), + ).toBe(false); + release!(); + await rejected; + expect(released).toBe(true); + expect(calls.some((call) => call.target === "s4")).toBe(false); + }); + it("bounds collection across many targets and stops scheduling after cancellation", async () => { const frames = Array.from({ length: 12 }, (_, i) => ({ frameId: `f${i}`, diff --git a/apps/extension/src/tools/vom/capture-coordinator.ts b/apps/extension/src/tools/vom/capture-coordinator.ts index 323c5a8..c8315c0 100644 --- a/apps/extension/src/tools/vom/capture-coordinator.ts +++ b/apps/extension/src/tools/vom/capture-coordinator.ts @@ -41,15 +41,25 @@ async function collectTargets( signal?: AbortSignal, ): Promise { let cursor = 0; + let failed = false; + let failure: unknown; await Promise.all( Array.from({ length: Math.min(4, items.length) }, async () => { - while (cursor < items.length) { - throwCaptureAborted(signal); - const item = items[cursor++]; - await collect(item); + try { + while (cursor < items.length && !failed && !signal?.aborted) { + const item = items[cursor++]; + await collect(item); + } + } catch (error) { + if (!failed) failure = error; + failed = true; } }), ); + // A worker's rejection must not let the capture outlive its own task scope. + // In-flight reads and their finally blocks settle before cancellation/failure. + throwCaptureAborted(signal); + if (failed) throw failure; } /** When topology discovery is unavailable, snapshot documents still carry @@ -91,6 +101,7 @@ export async function captureObservationFacts( } catch (error) { if (isCaptureAbort(error)) throw error; } + throwCaptureAborted(signal); const geometry = new GeometryContext(cdp, tabId, graph, signal); const issues: CaptureIssue[] = []; const groups = new Map>(); @@ -144,19 +155,19 @@ export async function captureObservationFacts( issues, signal, ); - if ( - !(await enrichFormControlStates( - scoped, - tabId, - batch.documents.map((doc) => doc.domNodes), - signal, - )) - ) - unavailable(batch, "forms"); + const formsAvailable = await enrichFormControlStates( + scoped, + tabId, + batch.documents.map((doc) => doc.domNodes), + signal, + ); + throwCaptureAborted(signal); + if (!formsAvailable) unavailable(batch, "forms"); const present = new Set(batch.documents.map((doc) => doc.frame.frameId)); for (const frame of batch.frames) if (!present.has(frame.frameId)) unavailable(batch, "dom", frame.frameId); } catch (error) { + throwCaptureAborted(signal); if (isCaptureAbort(error)) throw error; firstFailure ??= error; unavailable(batch, "dom"); @@ -168,6 +179,7 @@ export async function captureObservationFacts( throwCaptureAborted(signal); await scoped.send(tabId, "Accessibility.enable", {}); } catch (error) { + throwCaptureAborted(signal); if (isCaptureAbort(error)) throw error; firstFailure ??= error; unavailable(batch, "ax"); @@ -181,8 +193,10 @@ export async function captureObservationFacts( "Accessibility.getFullAXTree", frame.frameId === "root" ? {} : { frameId: frame.frameId }, ); + throwCaptureAborted(signal); batch.ax.push({ frame, nodes: result.nodes ?? [] }); } catch (error) { + throwCaptureAborted(signal); if (isCaptureAbort(error)) throw error; firstFailure ??= error; unavailable(batch, "ax", frame.frameId); diff --git a/apps/extension/src/tools/vom/form-capture.ts b/apps/extension/src/tools/vom/form-capture.ts index 8875680..826bc15 100644 --- a/apps/extension/src/tools/vom/form-capture.ts +++ b/apps/extension/src/tools/vom/form-capture.ts @@ -150,5 +150,6 @@ export async function enrichFormControlStates( return false; } finally { await cdp.send(tabId, "Runtime.releaseObjectGroup", { objectGroup }).catch(() => undefined); + throwIfAborted(signal); } }