From ecb18449be712e0a8bd4bb9d7ddddc2906dcac98 Mon Sep 17 00:00:00 2001 From: Abe Jellinek <1770299+AbeJellinek@users.noreply.github.com> Date: Wed, 6 May 2026 13:16:25 -0400 Subject: [PATCH] Support streaming SDT segments --- .../content/zotero/xpcom/pdfWorker/manager.js | 141 ++++++++++++++---- chrome/content/zotero/xpcom/reader.js | 34 +++++ 2 files changed, 146 insertions(+), 29 deletions(-) diff --git a/chrome/content/zotero/xpcom/pdfWorker/manager.js b/chrome/content/zotero/xpcom/pdfWorker/manager.js index dd423f65eb..daa69793c4 100644 --- a/chrome/content/zotero/xpcom/pdfWorker/manager.js +++ b/chrome/content/zotero/xpcom/pdfWorker/manager.js @@ -86,6 +86,23 @@ class PDFWorker { }); } + // Like _query, but onPartial fires for each isPartial message before the + // terminal response resolves the promise. Returns { promise, abort }. + _streamingQuery(action, data, transfer, onPartial) { + this._lastPromiseID++; + let id = this._lastPromiseID; + let promise = new Promise((resolve, reject) => { + this._waitingPromises[id] = { resolve, reject, onPartial }; + this._worker.postMessage({ id, action, data }, transfer); + }); + let abort = () => { + if (this._worker && this._waitingPromises[id]) { + this._worker.postMessage({ action: 'abort', id }); + } + }; + return { promise, abort }; + } + _init() { if (this._worker) return; this._worker = new Worker(WORKER_URL); @@ -93,11 +110,22 @@ class PDFWorker { let message = event.data; if ('responseID' in message) { let promise = this._waitingPromises[message.responseID]; - delete this._waitingPromises[message.responseID]; if (!promise) { Zotero.debug(`Received response from PDF worker for unknown request ${message.responseID}`); return; } + if (message.isPartial) { + if (promise.onPartial) { + try { + promise.onPartial(message.data); + } + catch (e) { + Zotero.logError(e); + } + } + return; + } + delete this._waitingPromises[message.responseID]; let { resolve, reject } = promise; if ('error' in message) { reject(new Error(JSON.stringify(message.error))); @@ -669,48 +697,103 @@ class PDFWorker { */ async getStructuredData(itemID, isPriority, password) { return this._enqueue(async () => { - let attachment = await Zotero.Items.getAsync(itemID); - + let prep = await this._prepareStructuredDataRequest(itemID); + if (!prep) return null; + let { attachment, contentType, buf } = prep; Zotero.debug(`Getting structured document text from item ${attachment.libraryKey}`); let t = new Date(); - - let contentType = attachment.attachmentContentType; - if (!(attachment.isPDFAttachment() - || attachment.isEPUBAttachment() - || attachment.isSnapshotAttachment())) { - throw new Error('Item must be a PDF, EPUB, or snapshot attachment'); - } - - let path = await attachment.getFilePathAsync(); - if (!path) { - return null; - } - let buf = await IOUtils.read(path); - buf = new Uint8Array(buf).buffer; - try { var result = await this._query('getStructuredDocumentText', { buf, contentType, password }, [buf]); } catch (e) { - let error = new Error(`Worker action 'getStructuredDocumentText' failed: ${JSON.stringify({ error: e.message })}`); - try { - error.name = JSON.parse(e.message).name; - } - catch (e) { - Zotero.logError(e); - } - Zotero.logError(error); - throw error; + throw this._wrapStructuredDataError(e); } - Zotero.debug(`Extracted structured document text for item ${attachment.libraryKey} in ${new Date() - t} ms`); - return result; }, isPriority); } + async _prepareStructuredDataRequest(itemID) { + let attachment = await Zotero.Items.getAsync(itemID); + if (!(attachment.isPDFAttachment() + || attachment.isEPUBAttachment() + || attachment.isSnapshotAttachment())) { + throw new Error('Item must be a PDF, EPUB, or snapshot attachment'); + } + let path = await attachment.getFilePathAsync(); + if (!path) return null; + let buf = await IOUtils.read(path); + return { + attachment, + contentType: attachment.attachmentContentType, + buf: new Uint8Array(buf).buffer, + }; + } + + _wrapStructuredDataError(e) { + let error = new Error(`Worker action 'getStructuredDocumentText' failed: ${JSON.stringify({ error: e.message })}`); + try { + error.name = JSON.parse(e.message).name; + } + catch (parseErr) { + Zotero.logError(parseErr); + } + Zotero.logError(error); + return error; + } + + // Streaming variant of getStructuredData. onChunk receives partial chunks + // ({ kind: 'partial', pages, content, pageIndexOffset, contentIndexOffset, + // pageIndexRange, totalPageCount }) followed by a final chunk + // ({ kind: 'final', structure }). Returns { promise, abort }. + getStructuredDataStream(itemID, onChunk, options = {}) { + let abortFn = null; + let aborted = false; + let { password, batchSize, isPriority } = options; + let promise = this._enqueue(async () => { + if (aborted) { + let e = new Error('Aborted'); + e.name = 'AbortError'; + throw e; + } + let prep = await this._prepareStructuredDataRequest(itemID); + if (!prep) return; + let { attachment, contentType, buf } = prep; + Zotero.debug(`Streaming structured document text from item ${attachment.libraryKey}`); + let t = new Date(); + try { + let { promise: queryPromise, abort } = this._streamingQuery( + 'getStructuredDocumentText', + { + buf, + contentType, + password, + streaming: true, + ...(batchSize ? { batchSize } : {}), + }, + [buf], + onChunk + ); + abortFn = abort; + if (aborted) abort(); + await queryPromise; + } + catch (e) { + throw this._wrapStructuredDataError(e); + } + Zotero.debug(`Streamed structured document text for item ${attachment.libraryKey} in ${new Date() - t} ms`); + }, isPriority); + return { + promise, + abort: () => { + aborted = true; + if (abortFn) abortFn(); + }, + }; + } + /** * Get data for recognizer-server * diff --git a/chrome/content/zotero/xpcom/reader.js b/chrome/content/zotero/xpcom/reader.js index 387c3d8358..4770dfbec1 100644 --- a/chrome/content/zotero/xpcom/reader.js +++ b/chrome/content/zotero/xpcom/reader.js @@ -644,6 +644,40 @@ class ReaderInstance { } }); }, + getSDTStream: (password, onChunk, onStart) => { + return new this._iframeWindow.Promise((resolve, reject) => { + let { promise, abort } = Zotero.PDFWorker.getStructuredDataStream( + this._item.id, + (chunk) => { + try { + onChunk(Cu.cloneInto(chunk, this._iframeWindow)); + } + catch (e) { + Zotero.logError(e); + } + }, + { + password, + isPriority: true, + } + ); + if (onStart) { + try { + onStart(Cu.exportFunction(() => abort(), this._iframeWindow)); + } + catch (e) { + Zotero.logError(e); + } + } + promise.then( + () => resolve(), + (e) => { + Zotero.logError(e); + reject(new this._iframeWindow.Error(e.message)); + } + ); + }); + }, }, this._iframeWindow, { cloneFunctions: true })); this._resolveInitPromise();