From 771d17adcadfb5dc7ebe394055bb89754a06ebd8 Mon Sep 17 00:00:00 2001 From: Dan Stillman Date: Mon, 20 Jul 2026 15:40:57 -0400 Subject: [PATCH] Guard best-match scoring against model changes A search could embed the query with one model's prefix on a worker initialized for another and compare it against vectors from a third. Tag the worker with the model version it was initialized for, make scoring wait out an in-progress model switch, refuse to score an index that wasn't stamped by the active model, and discard results if the model changes mid-scoring. The items view treats a not-ready index as an empty result rather than showing an unranked scope. --- .../content/zotero/collectionViewItemTree.jsx | 7 +- chrome/content/zotero/xpcom/embeddings.js | 65 +++++++++++++++++++ test/tests/embeddingsTest.js | 16 +++++ 3 files changed, 87 insertions(+), 1 deletion(-) diff --git a/chrome/content/zotero/collectionViewItemTree.jsx b/chrome/content/zotero/collectionViewItemTree.jsx index 8aaa78a771..880e2f2983 100644 --- a/chrome/content/zotero/collectionViewItemTree.jsx +++ b/chrome/content/zotero/collectionViewItemTree.jsx @@ -202,7 +202,12 @@ class CollectionViewItemTreeRowProvider extends ItemTreeRowProvider { // Scoring can fail while the model is still downloading or the // index is being rebuilt -- show no results rather than an // unranked scope - Zotero.logError(e); + if (e instanceof Zotero.Embeddings.IndexNotReadyError) { + Zotero.debug("Embeddings: index not ready for best-match search -- showing no results"); + } + else { + Zotero.logError(e); + } this._bestMatchRanks = new Map(); return []; } diff --git a/chrome/content/zotero/xpcom/embeddings.js b/chrome/content/zotero/xpcom/embeddings.js index 72b4f080fd..6d14bef9d9 100644 --- a/chrome/content/zotero/xpcom/embeddings.js +++ b/chrome/content/zotero/xpcom/embeddings.js @@ -456,6 +456,25 @@ Zotero.Embeddings = new function () { let _workerReady = null; let _requestID = 0; let _pending = new Map(); + // Model identity the current worker was initialized with + let _workerModelVersion = null; + // Bumped on every engine shutdown (e.g. a model switch), so long-running + // consumers can detect that the model changed under them and discard + // their results + let _modelGeneration = 0; + + /** + * Thrown when the stored embeddings can't be searched for the active + * model -- during a model switch, or while the index is being rebuilt + * after a revision bump. Callers should treat the index as still being + * prepared rather than scoring mismatched data. + */ + this.IndexNotReadyError = class extends Error { + constructor(message) { + super(message); + this.name = 'EmbeddingsIndexNotReadyError'; + } + }; function _ensureWorker() { if (_worker) { @@ -519,8 +538,15 @@ Zotero.Embeddings = new function () { // from the data directory) and the bundled ORT wasm binary. The model must // already be downloaded. async function _getWorker() { + // A worker initialized for a different model or revision can't be + // reused -- its weights and prefixes wouldn't match the active model + if (_worker && _workerModelVersion + && _workerModelVersion !== Zotero.Embeddings.getModelVersion()) { + Zotero.Embeddings.shutdownEngine(); + } _ensureWorker(); if (!_workerReady) { + let modelVersion = Zotero.Embeddings.getModelVersion(); _workerReady = (async () => { let model = _getModel(); if (!(await Zotero.Embeddings.isDownloaded())) { @@ -554,6 +580,7 @@ Zotero.Embeddings = new function () { wasmPaths: RESOURCE_DIR, wasmBinary: wasm.response }, transfer); + _workerModelVersion = modelVersion; Zotero.debug('Embeddings: worker initialized'); })(); } @@ -588,6 +615,8 @@ Zotero.Embeddings = new function () { _worker = null; } _workerReady = null; + _workerModelVersion = null; + _modelGeneration++; _failPending(new Error('Embeddings worker shut down')); }; @@ -695,7 +724,24 @@ Zotero.Embeddings = new function () { if (!itemIDs.length || !this.isEnabled()) { return scores; } + // Wait out any in-progress model switch, so the query isn't embedded + // with one model and compared against another's vectors + await Zotero.Embeddings.Indexing.waitForPendingModelSwitch(); await this.initDB(); + // The stored vectors must have been produced by the active model. + // During a switch, or a reindex after a revision bump, the database + // isn't stamped for the new model until the indexer starts filling it. + let modelVersion = this.getModelVersion(); + let indexedVersion = await Zotero.DB.valueQueryAsync( + "SELECT value FROM embeddings.itemEmbeddingsMeta WHERE key='modelVersion'" + ); + if (indexedVersion !== modelVersion) { + throw new this.IndexNotReadyError( + `Embeddings index is for '${indexedVersion || 'no model'}', ` + + `but the active model is '${modelVersion}'` + ); + } + let generation = _modelGeneration; let query = await this.embedQuery(queryText); let dim = query.length; @@ -703,6 +749,11 @@ Zotero.Embeddings = new function () { // parameter limit for large collections), scoring each as we go. let chunkSize = 500; for (let i = 0; i < itemIDs.length; i += chunkSize) { + // If the model changed while we were scoring, the scores computed + // so far mix models -- discard them + if (generation !== _modelGeneration) { + throw new this.IndexNotReadyError('Model changed during scoring'); + } let chunk = itemIDs.slice(i, i + chunkSize); let rows = await Zotero.DB.queryAsync( "SELECT itemID, embedding FROM embeddings.itemEmbeddings WHERE itemID IN (" @@ -718,6 +769,9 @@ Zotero.Embeddings = new function () { scores.set(row.itemID, dot); } } + if (generation !== _modelGeneration) { + throw new this.IndexNotReadyError('Model changed during scoring'); + } return scores; }; @@ -848,6 +902,17 @@ Zotero.Embeddings.Indexing = new function () { return _switchChain; } + /** + * Resolves once any in-progress model switch (stopping the indexer, + * clearing the old vectors, starting reindexing) has finished, so that + * callers don't operate across a switch + * + * @return {Promise} + */ + this.waitForPendingModelSwitch = function () { + return _switchChain; + }; + async function _doSwitchModel() { // Stop any in-progress indexing and wait for it to actually finish before // touching the stored vectors, so we never mix models/dimensions. diff --git a/test/tests/embeddingsTest.js b/test/tests/embeddingsTest.js index e45783681c..2f8d834f1a 100644 --- a/test/tests/embeddingsTest.js +++ b/test/tests/embeddingsTest.js @@ -24,6 +24,22 @@ describe("Zotero.Embeddings", function () { }); }); + describe("#scoreItemIDs()", function () { + it("should report the index as not ready when it wasn't built by the active model", async function () { + let stubs = [ + sinon.stub(Zotero.Embeddings, 'isEnabled').returns(true), + sinon.stub(Zotero.Embeddings, 'getModelVersion').returns('test-model/1') + ]; + try { + let e = await getPromiseError(Zotero.Embeddings.scoreItemIDs('query', [1])); + assert.instanceOf(e, Zotero.Embeddings.IndexNotReadyError); + } + finally { + stubs.forEach(stub => stub.restore()); + } + }); + }); + describe("#getScoreFraction()", function () { it("should clamp scores into the active model's display range", function () { // bge-small-en-v1.5's displayScoreRange is [0.5, 0.75]