diff --git a/sdk/typescript/src/knowledge-base.ts b/sdk/typescript/src/knowledge-base.ts index b7c26ea044..4b0bebddb6 100644 --- a/sdk/typescript/src/knowledge-base.ts +++ b/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, }); diff --git a/sdk/typescript/tests-ts/knowledge-base-pdf-cancellation.test.ts b/sdk/typescript/tests-ts/knowledge-base-pdf-cancellation.test.ts new file mode 100644 index 0000000000..3b68f2fff5 --- /dev/null +++ b/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([]); + } +});