fix(observation): await capture task cleanup before propagating cancellation

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<string, { resolve: () => void; reject: (error: Error) => void }>();
+    const calls: Array<{ target: string; method: string }> = [];
+    let release: (() => void) | undefined;
+    let released = false;
+    const send: NonNullable<CdpRunner["sendToTarget"]> = async <T>(
+      target: CdpTarget,
+      method: string,
+      params?: object,
+    ) => {
+      calls.push({ target: target.sessionId ?? "main", method });
+      if (method === "DOMSnapshot.captureSnapshot" && target.sessionId) {
+        await new Promise<void>((resolve, reject) =>
+          pending.set(target.sessionId!, { resolve, reject }),
+        );
+      }
+      if (method === "Runtime.releaseObjectGroup") {
+        await new Promise<void>((resolve) => {
+          release = resolve;
+        });
+        released = true;
+      }
+      const result = await original<T>(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}`,

apps/extension/src/tools/vom/capture-coordinator.ts

@@ -41,15 +41,25 @@ async function collectTargets<T>(
   signal?: AbortSignal,
 ): Promise<void> {
   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<T extends FrameOwnedAxNode>(
   } catch (error) {
     if (isCaptureAbort(error)) throw error;
   }
+  throwCaptureAborted(signal);
   const geometry = new GeometryContext(cdp, tabId, graph, signal);
   const issues: CaptureIssue[] = [];
   const groups = new Map<string, TargetBatch<T>>();
@@ -144,19 +155,19 @@ export async function captureObservationFacts<T extends FrameOwnedAxNode>(
           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<T extends FrameOwnedAxNode>(
         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<T extends FrameOwnedAxNode>(
             "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);

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