fix(knowledge): honor cancellation during PDF extraction

sdk/typescript/src/knowledge-base.ts

@@ -271,10 +271,13 @@ async function extractPdf(
     });
     try {
       const document = await loadingTask.promise;
+      signal?.throwIfAborted();
       const pages: string[] = [];
       for (let number = 1; number <= document.numPages; number++) {
         signal?.throwIfAborted();
-        const content = await (await document.getPage(number)).getTextContent();
+        const page = await document.getPage(number);
+        signal?.throwIfAborted();
+        const content = await page.getTextContent();
         signal?.throwIfAborted();
         pages.push(
           content.items
@@ -285,9 +288,10 @@ async function extractPdf(
       return pages.join("\n");
     } finally {
       await loadingTask.destroy();
+      signal?.throwIfAborted();
     }
   } catch (error) {
-    if (signal?.aborted) throw error;
+    signal?.throwIfAborted();
     throw new Error(`Cannot extract text from knowledge base PDF: ${path}`, {
       cause: error,
     });

sdk/typescript/tests-ts/knowledge-base-pdf-cancellation.test.ts

@@ -0,0 +1,133 @@
+import { mkdir, readFile, readdir, writeFile } from "node:fs/promises";
+import { join } from "node:path";
+import { afterEach, expect, mock, test } from "bun:test";
+import { createApiTestFixtures } from "./support/temporary-directories.js";
+import { runTestInSubprocess } from "./support/test-subprocess.js";
+
+const { temporaryDirectory, cleanup } = createApiTestFixtures();
+afterEach(cleanup);
+
+test.each([
+  "load",
+  "page",
+  "empty final page",
+  "during cleanup",
+  "cleanup failure",
+  "uncanceled cleanup failure",
+  "not canceled",
+] as const)("PDF extraction cancellation boundary: %s", async (boundary) => {
+  if (
+    runTestInSubprocess(
+      import.meta.path,
+      `PDF extraction cancellation boundary: ${boundary}`,
+    )
+  )
+    return;
+
+  const root = await temporaryDirectory();
+  const documentPath = join(root, "architecture.pdf");
+  const staging = join(root, "staging");
+  await mkdir(staging);
+  await writeFile(documentPath, "%PDF-synthetic-loader-fixture");
+  const reason = new Error("Synthetic PDF extraction cancellation");
+  const cleanupError = new Error("Synthetic PDF cleanup failure");
+  const emptyPage =
+    boundary === "empty final page" || boundary === "during cleanup";
+  const allPages =
+    boundary === "not canceled" || boundary === "uncanceled cleanup failure";
+  let controller = new AbortController();
+  let pages: number[] = [];
+  let textPages: number[] = [];
+  let destroyed = 0;
+  mock.module("pdfjs-dist/legacy/build/pdf.mjs", () => ({
+    VerbosityLevel: { ERRORS: 0 },
+    getDocument: () => ({
+      promise: Promise.resolve().then(() => {
+        if (boundary === "load") controller.abort(reason);
+        return {
+          numPages: emptyPage ? 1 : 3,
+          async getPage(number: number) {
+            pages.push(number);
+            if (boundary === "page") controller.abort(reason);
+            return {
+              async getTextContent() {
+                textPages.push(number);
+                if (
+                  boundary === "empty final page" ||
+                  boundary === "cleanup failure"
+                )
+                  controller.abort(reason);
+                return {
+                  items: emptyPage ? [] : [{ str: `Synthetic page ${number}` }],
+                };
+              },
+            };
+          },
+        };
+      }),
+      async destroy() {
+        destroyed++;
+        if (boundary === "during cleanup") controller.abort(reason);
+        if (
+          boundary === "cleanup failure" ||
+          boundary === "uncanceled cleanup failure"
+        )
+          throw cleanupError;
+      },
+    }),
+  }));
+  const { readKnowledgeBaseSnapshot, prepareKnowledgeBase } =
+    await import("../src/knowledge-base.js");
+  for (const entrypoint of ["snapshot", "prepared"] as const) {
+    controller = new AbortController();
+    pages = [];
+    textPages = [];
+    destroyed = 0;
+    const read = async () => {
+      if (entrypoint === "snapshot") {
+        const snapshot = await readKnowledgeBaseSnapshot(
+          [documentPath],
+          controller.signal,
+        );
+        return Object.values(snapshot.documents);
+      }
+      const prepared = await prepareKnowledgeBase(
+        [documentPath],
+        controller.signal,
+        staging,
+      );
+      try {
+        return await Promise.all(
+          (await readdir(prepared.path)).map((name) =>
+            readFile(join(prepared.path, name), "utf8"),
+          ),
+        );
+      } finally {
+        await prepared.cleanup();
+      }
+    };
+    if (boundary === "not canceled") {
+      await expect(read()).resolves.toEqual([
+        "Synthetic page 1\nSynthetic page 2\nSynthetic page 3",
+      ]);
+    } else if (boundary === "uncanceled cleanup failure") {
+      await expect(read()).rejects.toMatchObject({
+        cause: { cause: cleanupError },
+      });
+    } else {
+      await expect(read()).rejects.toBe(reason);
+    }
+    expect(pages).toEqual(
+      allPages ? [1, 2, 3] : boundary === "load" ? [] : [1],
+    );
+    expect(textPages).toEqual(
+      allPages
+        ? [1, 2, 3]
+        : boundary === "load" || boundary === "page"
+          ? []
+          : [1],
+    );
+    expect(destroyed).toBe(1);
+    expect(await readdir(staging)).toEqual([]);
+  }
+});