import { createHash } from 'node:crypto' import { Worker } from 'node:worker_threads' import type { KnowledgeChunkingSettings } from '../shared/knowledge-contracts' import { assertDocumentBuffer, chunkDocumentAdvanced, DocumentTextUnavailableError, extractPdfTextPages, parseDocument, type DocumentChunk, type ParsedDocument, type PdfTextExtraction, type PdfTextExtractionOptions } from './knowledge/document-parser' import { extractPptxPages, type PptxPage } from './knowledge/pptx-parser' import type { DocumentParseOperation, DocumentParseRequest, DocumentParseResponse, SerializedDocumentParseError } from './document-parse-errors' export type { DocumentParseOperation } from './document-parse-errors' // PERF-15/16: pure document work (PDF / Office / text parsing, chunking, // hashing) runs in a single worker thread with a FIFO queue, so Main only // orchestrates. Database writes stay on the caller's side, in the same order. // Same fallback rules as readonly-query-reader: if the worker cannot start or // crashes, the call runs inline (the previous behaviour) and the worker is // retried after a backoff. Aborting a running task terminates the worker. /** Worker infrastructure failure; callers fall back to the inline path. */ export class DocumentParseWorkerUnavailableError extends Error { constructor(message: string, options?: { cause?: unknown }) { super(message, options) this.name = 'DocumentParseWorkerUnavailableError' } } type Task = { id: number op: DocumentParseOperation args: unknown[] transfer: ArrayBuffer[] resolve: (value: unknown) => void reject: (error: unknown) => void cleanup: () => void } const RESTART_BACKOFF_MS = 30_000 type PdfExceptionConstructor = new (message: string) => Error type PdfExceptionClasses = Partial> const pdfExceptionNames = [ 'AbortException', 'InvalidPDFException', 'PasswordException', 'ResponseException', 'UnknownErrorException' ] as const let pdfExceptions: Promise | undefined /** PDF.js exception classes, loaded only once a PDF error crosses the boundary. */ function loadPdfExceptions(): Promise { pdfExceptions ??= import('pdfjs-dist/legacy/build/pdf.mjs').then( (pdfjs) => { const module = pdfjs as unknown as Record return Object.fromEntries(pdfExceptionNames.flatMap((name) => typeof module[name] === 'function' ? [[name, module[name] as PdfExceptionConstructor]] : [])) }, () => ({}) ) return pdfExceptions } /** * Rebuilds an error with the same class, name, message, cause and primitive * fields as the one thrown inline (stack traces differ by nature). */ export async function deserializeDocumentParseError(error: SerializedDocumentParseError): Promise { let result: Error if (error.name === 'DocumentTextUnavailableError') { result = new DocumentTextUnavailableError(error.message) } else if ((pdfExceptionNames as readonly string[]).includes(error.name)) { const Constructor = (await loadPdfExceptions())[error.name] result = Constructor ? new Constructor(error.message) : new Error(error.message) } else { const Constructor = error.name === 'RangeError' ? RangeError : error.name === 'TypeError' ? TypeError : Error result = new Constructor(error.message) } if (result.name !== error.name) result.name = error.name if (error.cause) { Object.defineProperty(result, 'cause', { value: await deserializeDocumentParseError(error.cause), writable: true, configurable: true }) } for (const [key, value] of Object.entries(error.extra ?? {})) { Object.defineProperty(result, key, { value, writable: true, configurable: true, enumerable: true }) } return result } export class DocumentParseWorkerClient { private worker?: Worker private nextId = 0 private disabledUntil = 0 private closed = false private active?: Task private readonly queue: Task[] = [] constructor( private readonly workerPath: string, private readonly now: () => number = Date.now ) {} /** False while closed or backing off after a worker failure. */ get available(): boolean { return !this.closed && this.now() >= this.disabledUntil } get pendingCount(): number { return this.queue.length + (this.active ? 1 : 0) } call( op: DocumentParseOperation, args: unknown[], transfer: ArrayBuffer[] = [], signal?: AbortSignal ): Promise { if (this.closed) return Promise.reject(new DocumentParseWorkerUnavailableError('Document parse worker closed')) if (signal?.aborted) return Promise.reject(signal.reason) if (this.now() < this.disabledUntil) { return Promise.reject(new DocumentParseWorkerUnavailableError('Document parse worker is backing off')) } return new Promise((resolve, reject) => { const task: Task = { id: ++this.nextId, op, args, transfer, resolve: resolve as (value: unknown) => void, reject, cleanup: () => signal?.removeEventListener('abort', onAbort) } const onAbort = (): void => { if (this.active === task) { // The task is CPU-bound in the worker; stop it by terminating the worker. this.active = undefined const worker = this.worker this.worker = undefined if (worker) void worker.terminate() } else { const index = this.queue.indexOf(task) if (index < 0) return this.queue.splice(index, 1) } task.cleanup() reject(signal!.reason) this.pump() } signal?.addEventListener('abort', onAbort, { once: true }) this.queue.push(task) this.pump() }) } private pump(): void { if (this.active || this.closed) return const task = this.queue.shift() if (!task) { this.worker?.unref() return } let worker: Worker try { worker = this.ensureWorker() } catch (error) { this.disabledUntil = this.now() + RESTART_BACKOFF_MS this.failAll(task, new DocumentParseWorkerUnavailableError('Document parse worker failed to start', { cause: error })) return } this.active = task // Keep the process alive only while a task is outstanding. worker.ref() try { worker.postMessage( { id: task.id, op: task.op, args: task.args } satisfies DocumentParseRequest, task.transfer ) } catch (error) { this.active = undefined task.cleanup() task.reject(new DocumentParseWorkerUnavailableError('Document parse task could not be posted', { cause: error })) this.pump() } } private failAll(first: Task | undefined, error: Error): void { const tasks = [...(first ? [first] : []), ...this.queue.splice(0)] for (const task of tasks) { task.cleanup() task.reject(error) } } private ensureWorker(): Worker { if (this.worker) return this.worker const worker = new Worker(this.workerPath) this.worker = worker worker.on('message', (message: DocumentParseResponse) => { const task = this.active if (this.worker !== worker || !task || task.id !== message.id) return this.active = undefined task.cleanup() if ('error' in message) { void deserializeDocumentParseError(message.error).then(task.reject) } else { task.resolve(message.result) } this.pump() }) const failed = (error: Error): void => { // Ignore exits of workers that were terminated on purpose (abort, close). if (this.worker !== worker) return this.worker = undefined this.disabledUntil = this.now() + RESTART_BACKOFF_MS const active = this.active this.active = undefined this.failAll(active, new DocumentParseWorkerUnavailableError(error.message, { cause: error })) void worker.terminate() } worker.once('error', failed) worker.once('exit', (code) => failed(new Error(`Document parse worker exited (${code})`))) worker.unref() return worker } /** Test hook: simulate a crash by terminating the worker. */ async terminateWorkerForTest(): Promise { await this.worker?.terminate() } /** Test hook: end the post-crash backoff. */ resetBackoffForTest(): void { this.disabledUntil = 0 } close(): void { this.closed = true const worker = this.worker this.worker = undefined const active = this.active this.active = undefined this.failAll(active, new DocumentParseWorkerUnavailableError('Document parse worker closed')) if (worker) void worker.terminate() } } let sharedClient: DocumentParseWorkerClient | undefined /** Enables the parse worker for this process (Main calls it once at startup). */ export function configureDocumentParseWorker(workerPath: string | undefined): void { sharedClient?.close() sharedClient = workerPath ? new DocumentParseWorkerClient(workerPath) : undefined } export function closeDocumentParseWorker(): void { sharedClient?.close() sharedClient = undefined } /** A private copy of the bytes that can be transferred without detaching the caller's buffer. */ function transferableCopy(buffer: Uint8Array): ArrayBuffer { return Uint8Array.prototype.slice.call(buffer).buffer as ArrayBuffer } /** * Runs `op` in the worker when available, otherwise inline. Worker * infrastructure failures fall back to the inline path; operation errors and * aborts propagate unchanged. An already-aborted signal also runs inline, so * validation and abort errors surface in exactly the inline order. */ export async function runOffMain( client: DocumentParseWorkerClient | undefined, op: DocumentParseOperation, request: () => { args: unknown[]; transfer?: ArrayBuffer[] }, signal: AbortSignal | undefined, inline: () => T | Promise ): Promise { if (!client?.available || signal?.aborted) return inline() const { args, transfer } = request() try { return await client.call(op, args, transfer, signal) } catch (error) { if (error instanceof DocumentParseWorkerUnavailableError && !signal?.aborted) return inline() throw error } } export function parseDocumentOffMain( name: string, buffer: Buffer, signal?: AbortSignal, client: DocumentParseWorkerClient | undefined = sharedClient ): Promise { // Cheap size checks stay inline so oversized input is never copied. try { assertDocumentBuffer(buffer) } catch (error) { return Promise.reject(error) } return runOffMain(client, 'parseDocument', () => { const bytes = transferableCopy(buffer) return { args: [name, bytes], transfer: [bytes] } }, signal, () => parseDocument(name, buffer, signal)) } export function extractPdfTextPagesOffMain( buffer: Buffer, options: PdfTextExtractionOptions = {}, client: DocumentParseWorkerClient | undefined = sharedClient ): Promise { const { signal, ...limits } = options return runOffMain(client, 'extractPdfTextPages', () => { const bytes = transferableCopy(buffer) return { args: [bytes, limits], transfer: [bytes] } }, signal, () => extractPdfTextPages(buffer, options)) } export function extractPptxPagesOffMain( buffer: Buffer, client: DocumentParseWorkerClient | undefined = sharedClient ): Promise { return runOffMain(client, 'extractPptxPages', () => { const bytes = transferableCopy(buffer) return { args: [bytes], transfer: [bytes] } }, undefined, () => extractPptxPages(buffer)) } /** * Not cancellable on purpose: the inline call it replaces was synchronous, so * an abort during chunking keeps surfacing at the caller's next checkpoint. */ export function chunkDocumentOffMain( document: ParsedDocument, settings: KnowledgeChunkingSettings, client: DocumentParseWorkerClient | undefined = sharedClient ): Promise { return runOffMain(client, 'chunkDocument', () => ({ // Only the fields chunking reads; images (Buffers) would otherwise be cloned. args: [{ title: document.title, sourceFormat: document.sourceFormat, content: '', sections: document.sections, warnings: [] }, settings] }), undefined, () => chunkDocumentAdvanced(document, settings)) } /** Hex SHA-256 of `buffer`, computed in the worker for large inputs. */ export function sha256OffMain( buffer: Uint8Array, client: DocumentParseWorkerClient | undefined = sharedClient ): Promise { const inline = (): string => createHash('sha256').update(buffer).digest('hex') // Small inputs hash faster than a round trip. if (buffer.byteLength < 1024 * 1024) return Promise.resolve().then(inline) return runOffMain(client, 'sha256', () => { const bytes = transferableCopy(buffer) return { args: [bytes], transfer: [bytes] } }, undefined, inline) }