Bug 1911795 - Only await on createEngine and don't rely on the progress callback to know when the engine is ready r=tarek,atossou

Differential Revision: https://phabricator.services.mozilla.com/D218836

toolkit/components/pdfjs/content/PdfjsParent.sys.mjs

@@ -18,6 +18,7 @@ import { XPCOMUtils } from "resource://gre/modules/XPCOMUtils.sys.mjs";
 const lazy = {};
 
 ChromeUtils.defineESModuleGetters(lazy, {
+  clearTimeout: "resource://gre/modules/Timer.sys.mjs",
   createEngine: "chrome://global/content/ml/EngineProcess.sys.mjs",
   EngineProcess: "chrome://global/content/ml/EngineProcess.sys.mjs",
   IndexedDBCache: "chrome://global/content/ml/ModelHub.sys.mjs",
@@ -27,10 +28,12 @@ ChromeUtils.defineESModuleGetters(lazy, {
   PdfJsTelemetry: "resource://pdf.js/PdfJsTelemetry.sys.mjs",
   PrivateBrowsingUtils: "resource://gre/modules/PrivateBrowsingUtils.sys.mjs",
   SetClipboardSearchString: "resource://gre/modules/Finder.sys.mjs",
+  setTimeout: "resource://gre/modules/Timer.sys.mjs",
 });
 
 const IMAGE_TO_TEXT_TASK = "moz-image-to-text";
 const ML_ENGINE_ID = "pdfjs";
+const ML_ENGINE_MAX_TIMEOUT = 60000;
 
 var Svc = {};
 XPCOMUtils.defineLazyServiceGetter(
@@ -160,12 +163,34 @@ export class PdfjsParent extends JSWindowActorParent {
     if (service !== IMAGE_TO_TEXT_TASK) {
       throw new Error("Invalid service");
     }
-    const { promise, resolve } = Promise.withResolvers();
-    const self = this;
-    const aggregator = new lazy.MultiProgressAggregator({
-      progressCallback({ ok, total, totalLoaded, statusText }) {
-        const finished = statusText === lazy.Progress.ProgressStatusText.DONE;
-        if (listenToProgress) {
+
+    let aggregator = null;
+    if (listenToProgress) {
+      const self = this;
+      const timeoutCallback = () => {
+        lazy.clearTimeout(timeoutId);
+        timeoutId = null;
+        self.sendAsyncMessage("PDFJS:Child:handleEvent", {
+          type: "loadAIEngineProgress",
+          detail: {
+            service,
+            ok: false,
+            finished: true,
+          },
+        });
+      };
+      let timeoutId = lazy.setTimeout(timeoutCallback, ML_ENGINE_MAX_TIMEOUT);
+      aggregator = new lazy.MultiProgressAggregator({
+        progressCallback({ ok, total, totalLoaded, statusText }) {
+          if (timeoutId !== null) {
+            lazy.clearTimeout(timeoutId);
+            timeoutId = lazy.setTimeout(timeoutCallback, ML_ENGINE_MAX_TIMEOUT);
+          } else {
+            // The timeout has already fired, so we don't need to do anything.
+            this.progressCallback = null;
+            return;
+          }
+          const finished = statusText === lazy.Progress.ProgressStatusText.DONE;
           self.sendAsyncMessage("PDFJS:Child:handleEvent", {
             type: "loadAIEngineProgress",
             detail: {
@@ -176,28 +201,23 @@ export class PdfjsParent extends JSWindowActorParent {
               finished,
             },
           });
-        }
-        if (finished) {
-          // Once we're done, we can remove the progress callback.
-          this.progressCallback = null;
-          resolve(ok);
-        }
-      },
-      watchedTypes: [
-        lazy.Progress.ProgressType.DOWNLOAD,
-        lazy.Progress.ProgressType.LOAD_FROM_CACHE,
-      ],
-    });
-
-    if (Cu.isInAutomation) {
-      return !!(await this.#createAIEngine("moz-eco", null));
+          if (finished) {
+            lazy.clearTimeout(timeoutId);
+            // Once we're done, we can remove the progress callback.
+            this.progressCallback = null;
+          }
+        },
+        watchedTypes: [
+          lazy.Progress.ProgressType.DOWNLOAD,
+          lazy.Progress.ProgressType.LOAD_FROM_CACHE,
+        ],
+      });
     }
 
-    const [engine, ok] = await Promise.all([
-      this.#createAIEngine(service, aggregator),
-      promise,
-    ]);
-    return !!engine && ok;
+    return !!(await this.#createAIEngine(
+      Cu.isInAutomation ? "moz-echo" : service,
+      aggregator
+    ));
   }
 
   async _mlDelete({ data: service }) {
@@ -210,8 +230,7 @@ export class PdfjsParent extends JSWindowActorParent {
       await lazy.EngineProcess.destroyMLEngine();
       const cache = await lazy.IndexedDBCache.init();
       await cache.deleteModels({
-        model: "mozilla/distilvit",
-        revision: "main",
+        taskName: service,
       });
     } catch (e) {
       console.error("Failed to delete AI model", e);

toolkit/components/pdfjs/test/browser.toml

@@ -7,6 +7,9 @@ support-files = [
 ["browser_pdfjs_alttext_load_engine.js"]
 skip-if = ["!nightly_build"] # Bug 1890946
 
+["browser_pdfjs_alttext_two_tabs.js"]
+skip-if = ["!nightly_build"] # Bug 1890946
+
 ["browser_pdfjs_caret_browsing_mode.js"]
 
 ["browser_pdfjs_chrome_events.js"]

toolkit/components/pdfjs/test/browser_pdfjs_alttext_two_tabs.js

@@ -0,0 +1,98 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+Services.scriptloader.loadSubScript(
+  "chrome://mochitests/content/browser/toolkit/components/ml/tests/browser/head.js",
+  this
+);
+
+const RELATIVE_DIR = "toolkit/components/pdfjs/test/";
+const TESTROOT = "https://example.com/browser/" + RELATIVE_DIR;
+const pdfUrl = TESTROOT + "file_pdfjs_test.pdf";
+const altTextPref = "pdfjs.enableAltText";
+const guessAltTextPref = "pdfjs.enableGuessAltText";
+const browserMLPref = "browser.ml.enable";
+const annotationEditorModePref = "pdfjs.annotationEditorMode";
+
+async function setupRemoteClient() {
+  const { removeMocks, remoteClients } = await createAndMockMLRemoteSettings({
+    autoDownloadFromRemoteSettings: false,
+  });
+
+  return {
+    remoteClients,
+    async cleanup() {
+      await removeMocks();
+      await waitForCondition(
+        () => EngineProcess.areAllEnginesTerminated(),
+        "Waiting for all of the engines to be terminated.",
+        100,
+        200
+      );
+    },
+  };
+}
+
+add_task(async function test_loaded_model_in_two_tabs() {
+  let mimeService = Cc["@mozilla.org/mime;1"].getService(Ci.nsIMIMEService);
+  let handlerInfo = mimeService.getFromTypeAndExtension(
+    "application/pdf",
+    "pdf"
+  );
+
+  // Make sure pdf.js is the default handler.
+  is(
+    handlerInfo.alwaysAskBeforeHandling,
+    false,
+    "pdf handler defaults to always-ask is false"
+  );
+  is(
+    handlerInfo.preferredAction,
+    Ci.nsIHandlerInfo.handleInternally,
+    "pdf handler defaults to internal"
+  );
+
+  info("Pref action: " + handlerInfo.preferredAction);
+
+  await SpecialPowers.pushPrefEnv({
+    set: [
+      [altTextPref, true],
+      [browserMLPref, true],
+      [guessAltTextPref, true],
+      [annotationEditorModePref, 0],
+    ],
+  });
+  const setRC = await setupRemoteClient();
+  const { EngineProcess } = ChromeUtils.importESModule(
+    "chrome://global/content/ml/EngineProcess.sys.mjs"
+  );
+
+  const browsers = [];
+  for (let i = 0; i < 2; i++) {
+    const tab = await BrowserTestUtils.openNewForegroundTab(
+      gBrowser,
+      "about:blank"
+    );
+    const browser = tab.linkedBrowser;
+    browsers.push(browser);
+
+    await waitForPdfJS(browser, pdfUrl);
+    await enableEditor(browser, "Stamp", 1);
+
+    const isEnabled = await SpecialPowers.spawn(browser, [], () => {
+      const viewer = content.wrappedJSObject.PDFViewerApplication;
+      return viewer.mlManager.isEnabledFor("altText");
+    });
+    ok(isEnabled, `AltText is enabled in tab number ${i + 1}`);
+  }
+
+  for (const browser of browsers) {
+    await waitForPdfJSClose(browser, /* closeTab = */ true);
+  }
+
+  await setRC.remoteClients["ml-onnx-runtime"].resolvePendingDownloads(1);
+  await EngineProcess.destroyMLEngine();
+  await setRC.cleanup();
+
+  await SpecialPowers.popPrefEnv();
+});