4 files changed, 508 insertions, 0 deletions
diff --git a/toolkit/components/ml/content/EngineProcess.sys.mjs b/toolkit/components/ml/content/EngineProcess.sys.mjs
new file mode 100644
index 0000000000..36a9381192
--- /dev/null
+++ b/toolkit/components/ml/content/EngineProcess.sys.mjs
@@ -0,0 +1,241 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+const lazy = {};
+ChromeUtils.defineESModuleGetters(lazy, {
+  HiddenFrame: "resource://gre/modules/HiddenFrame.sys.mjs",
+});
+
+/**
+ * @typedef {import("../actors/MLEngineParent.sys.mjs").MLEngineParent} MLEngineParent
+ */
+
+/**
+ * @typedef {import("../../translations/actors/TranslationsEngineParent.sys.mjs").TranslationsEngineParent} TranslationsEngineParent
+ */
+
+/**
+ * This class controls the life cycle of the engine process used both in the
+ * Translations engine and the MLEngine component.
+ */
+export class EngineProcess {
+  /**
+   * @type {Promise<{ hiddenFrame: HiddenFrame, actor: TranslationsEngineParent }> | null}
+   */
+
+  /** @type {Promise<HiddenFrame> | null} */
+  static #hiddenFrame = null;
+  /** @type {Promise<TranslationsEngineParent> | null} */
+  static translationsEngineParent = null;
+  /** @type {Promise<MLEngineParent> | null} */
+  static mlEngineParent = null;
+
+  /** @type {((actor: TranslationsEngineParent) => void) | null} */
+  resolveTranslationsEngineParent = null;
+
+  /** @type {((actor: MLEngineParent) => void) | null} */
+  resolveMLEngineParent = null;
+
+  /**
+   * See if all engines are terminated. This is useful for testing.
+   *
+   * @returns {boolean}
+   */
+  static areAllEnginesTerminated() {
+    return (
+      !EngineProcess.#hiddenFrame &&
+      !EngineProcess.translationsEngineParent &&
+      !EngineProcess.mlEngineParent
+    );
+  }
+
+  /**
+   * @returns {Promise<TranslationsEngineParent>}
+   */
+  static async getTranslationsEngineParent() {
+    if (!this.translationsEngineParent) {
+      this.translationsEngineParent = this.#attachBrowser({
+        id: "translations-engine-browser",
+        url: "chrome://global/content/translations/translations-engine.html",
+        resolverName: "resolveTranslationsEngineParent",
+      });
+    }
+    return this.translationsEngineParent;
+  }
+
+  /**
+   * @returns {Promise<MLEngineParent>}
+   */
+  static async getMLEngineParent() {
+    if (!this.mlEngineParent) {
+      this.mlEngineParent = this.#attachBrowser({
+        id: "ml-engine-browser",
+        url: "chrome://global/content/ml/MLEngine.html",
+        resolverName: "resolveMLEngineParent",
+      });
+    }
+    return this.mlEngineParent;
+  }
+
+  /**
+   * @param {object} config
+   * @param {string} config.url
+   * @param {string} config.id
+   * @param {string} config.resolverName
+   * @returns {Promise<TranslationsEngineParent>}
+   */
+  static async #attachBrowser({ url, id, resolverName }) {
+    const hiddenFrame = await this.#getHiddenFrame();
+    const chromeWindow = await hiddenFrame.get();
+    const doc = chromeWindow.document;
+
+    if (doc.getElementById(id)) {
+      throw new Error(
+        "Attempting to append the translations-engine.html <browser> when one " +
+          "already exists."
+      );
+    }
+
+    const browser = doc.createXULElement("browser");
+    browser.setAttribute("id", id);
+    browser.setAttribute("remote", "true");
+    browser.setAttribute("remoteType", "web");
+    browser.setAttribute("disableglobalhistory", "true");
+    browser.setAttribute("type", "content");
+    browser.setAttribute("src", url);
+
+    ChromeUtils.addProfilerMarker(
+      "EngineProcess",
+      {},
+      `Creating the "${id}" process`
+    );
+    doc.documentElement.appendChild(browser);
+
+    const { promise, resolve } = Promise.withResolvers();
+
+    // The engine parents must resolve themselves when they are ready.
+    this[resolverName] = resolve;
+
+    return promise;
+  }
+
+  /**
+   * @returns {HiddenFrame}
+   */
+  static async #getHiddenFrame() {
+    if (!EngineProcess.#hiddenFrame) {
+      EngineProcess.#hiddenFrame = new lazy.HiddenFrame();
+    }
+    return EngineProcess.#hiddenFrame;
+  }
+
+  /**
+   * Destroy the translations engine, and remove the hidden frame if no other
+   * engines exist.
+   */
+  static destroyTranslationsEngine() {
+    return this.#destroyEngine({
+      id: "translations-engine-browser",
+      keyName: "translationsEngineParent",
+    });
+  }
+
+  /**
+   * Destroy the ML engine, and remove the hidden frame if no other engines exist.
+   */
+  static destroyMLEngine() {
+    return this.#destroyEngine({
+      id: "ml-engine-browser",
+      keyName: "mlEngineParent",
+    });
+  }
+
+  /**
+   * Destroy the specified engine and maybe the entire hidden frame as well if no engines
+   * are remaining.
+   */
+  static #destroyEngine({ id, keyName }) {
+    ChromeUtils.addProfilerMarker(
+      "EngineProcess",
+      {},
+      `Destroying the "${id}" engine`
+    );
+
+    const actorShutdown = this.forceActorShutdown(id, keyName).catch(
+      error => void console.error(error)
+    );
+
+    this[keyName] = null;
+
+    const hiddenFrame = EngineProcess.#hiddenFrame;
+    if (hiddenFrame && !this.translationsEngineParent && !this.mlEngineParent) {
+      EngineProcess.#hiddenFrame = null;
+
+      // Both actors are destroyed, also destroy the hidden frame.
+      actorShutdown.then(() => {
+        // Double check a race condition that no new actors have been created during
+        // shutdown.
+        if (this.translationsEngineParent && this.mlEngineParent) {
+          return;
+        }
+        if (!hiddenFrame) {
+          return;
+        }
+        hiddenFrame.destroy();
+        ChromeUtils.addProfilerMarker(
+          "EngineProcess",
+          {},
+          `Removing the hidden frame`
+        );
+      });
+    }
+
+    // Infallibly resolve the promise even if there are errors.
+    return Promise.resolve();
+  }
+
+  /**
+   * Shut down an actor and remove its <browser> element.
+   *
+   * @param {string} id
+   * @param {string} keyName
+   */
+  static async forceActorShutdown(id, keyName) {
+    const actorPromise = this[keyName];
+    if (!actorPromise) {
+      return;
+    }
+
+    let actor;
+    try {
+      actor = await actorPromise;
+    } catch {
+      // The actor failed to initialize, so it doesn't need to be shut down.
+      return;
+    }
+
+    // Shut down the actor.
+    try {
+      await actor.forceShutdown();
+    } catch (error) {
+      console.error("Failed to shut down the actor " + id, error);
+      return;
+    }
+
+    if (!EngineProcess.#hiddenFrame) {
+      // The hidden frame was already removed.
+      return;
+    }
+
+    // Remove the <brower> element.
+    const chromeWindow = EngineProcess.#hiddenFrame.getWindow();
+    const doc = chromeWindow.document;
+    const element = doc.getElementById(id);
+    if (!element) {
+      console.error("Could not find the <browser> element for " + id);
+      return;
+    }
+    element.remove();
+  }
+}
diff --git a/toolkit/components/ml/content/MLEngine.html b/toolkit/components/ml/content/MLEngine.html
new file mode 100644
index 0000000000..8763995102
--- /dev/null
+++ b/toolkit/components/ml/content/MLEngine.html
@@ -0,0 +1,16 @@
+<!-- This Source Code Form is subject to the terms of the Mozilla Public
+   - License, v. 2.0. If a copy of the MPL was not distributed with this
+   - file, You can obtain one at http://mozilla.org/MPL/2.0/. -->
+
+<!DOCTYPE html>
+<html>
+  <head>
+    <meta charset="utf-8" />
+    <meta
+      http-equiv="Content-Security-Policy"
+      content="default-src chrome: resource:; object-src 'none'"
+    />
+    <!-- Run the machine learning inference engine in its own singleton content process. -->
+  </head>
+  <body></body>
+</html>
diff --git a/toolkit/components/ml/content/MLEngine.worker.mjs b/toolkit/components/ml/content/MLEngine.worker.mjs
new file mode 100644
index 0000000000..1013977e07
--- /dev/null
+++ b/toolkit/components/ml/content/MLEngine.worker.mjs
@@ -0,0 +1,91 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+import { PromiseWorker } from "resource://gre/modules/workers/PromiseWorker.mjs";
+
+// Respect the preference "browser.ml.logLevel".
+let _loggingLevel = "Error";
+function log(...args) {
+  if (_loggingLevel !== "Error" && _loggingLevel !== "Warn") {
+    console.log("ML:", ...args);
+  }
+}
+function trace(...args) {
+  if (_loggingLevel === "Trace" || _loggingLevel === "All") {
+    console.log("ML:", ...args);
+  }
+}
+
+/**
+ * The actual MLEngine lives here in a worker.
+ */
+class MLEngineWorker {
+  /** @type {ArrayBuffer} */
+  #wasm;
+  /** @type {ArrayBuffer} */
+  #model;
+
+  constructor() {
+    // Connect the provider to the worker.
+    this.#connectToPromiseWorker();
+  }
+
+  /**
+   * @param {ArrayBuffer} wasm
+   * @param {ArrayBuffer} model
+   * @param {string} loggingLevel
+   */
+  initializeEngine(wasm, model, loggingLevel) {
+    this.#wasm = wasm;
+    this.#model = model;
+    _loggingLevel = loggingLevel;
+    // TODO - Initialize the engine for real here.
+    log("MLEngineWorker is initalized");
+  }
+
+  /**
+   * Run the worker.
+   *
+   * @param {string} request
+   */
+  run(request) {
+    if (!this.#wasm) {
+      throw new Error("Expected the wasm to exist.");
+    }
+    if (!this.#model) {
+      throw new Error("Expected the model to exist");
+    }
+    if (request === "throw") {
+      throw new Error(
+        'Received the message "throw", so intentionally throwing an error.'
+      );
+    }
+    trace("inference run requested with:", request);
+    return request.slice(0, Math.floor(request.length / 2));
+  }
+
+  /**
+   * Glue code to connect the `MLEngineWorker` to the PromiseWorker interface.
+   */
+  #connectToPromiseWorker() {
+    const worker = new PromiseWorker.AbstractWorker();
+    worker.dispatch = (method, args = []) => {
+      if (!this[method]) {
+        throw new Error("Method does not exist: " + method);
+      }
+      return this[method](...args);
+    };
+    worker.close = () => self.close();
+    worker.postMessage = (message, ...transfers) => {
+      self.postMessage(message, ...transfers);
+    };
+
+    self.addEventListener("message", msg => worker.handleMessage(msg));
+    self.addEventListener("unhandledrejection", function (error) {
+      throw error.reason;
+    });
+  }
+}
+
+new MLEngineWorker();
diff --git a/toolkit/components/ml/content/SummarizerModel.sys.mjs b/toolkit/components/ml/content/SummarizerModel.sys.mjs
new file mode 100644
index 0000000000..7cac55d92f
--- /dev/null
+++ b/toolkit/components/ml/content/SummarizerModel.sys.mjs
@@ -0,0 +1,160 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+/**
+ * @typedef {object} LazyImports
+ * @property {typeof import("../actors/MLEngineParent.sys.mjs").MLEngineParent} MLEngineParent
+ */
+
+/** @type {LazyImports} */
+const lazy = {};
+
+ChromeUtils.defineESModuleGetters(lazy, {
+  RemoteSettings: "resource://services-settings/remote-settings.sys.mjs",
+  TranslationsParent: "resource://gre/actors/TranslationsParent.sys.mjs",
+});
+
+ChromeUtils.defineLazyGetter(lazy, "console", () => {
+  return console.createInstance({
+    maxLogLevelPref: "browser.ml.logLevel",
+    prefix: "ML",
+  });
+});
+
+export class SummarizerModel {
+  /**
+   * The RemoteSettingsClient that downloads the summarizer model.
+   *
+   * @type {RemoteSettingsClient | null}
+   */
+  static #remoteClient = null;
+
+  /** @type {Promise<WasmRecord> | null} */
+  static #modelRecord = null;
+
+  /**
+   * The following constant controls the major version for wasm downloaded from
+   * Remote Settings. When a breaking change is introduced, Nightly will have these
+   * numbers incremented by one, but Beta and Release will still be on the previous
+   * version. Remote Settings will ship both versions of the records, and the latest
+   * asset released in that version will be used. For instance, with a major version
+   * of "1", assets can be downloaded for "1.0", "1.2", "1.3beta", but assets marked
+   * as "2.0", "2.1", etc will not be downloaded.
+   */
+  static MODEL_MAJOR_VERSION = 1;
+
+  /**
+   * Remote settings isn't available in tests, so provide mocked responses.
+   */
+  static mockRemoteSettings(remoteClient) {
+    lazy.console.log("Mocking remote client in SummarizerModel.");
+    SummarizerModel.#remoteClient = remoteClient;
+    SummarizerModel.#modelRecord = null;
+  }
+
+  /**
+   * Remove anything that could have been mocked.
+   */
+  static removeMocks() {
+    lazy.console.log("Removing mocked remote client in SummarizerModel.");
+    SummarizerModel.#remoteClient = null;
+    SummarizerModel.#modelRecord = null;
+  }
+  /**
+   * Download or load the model from remote settings.
+   *
+   * @returns {Promise<ArrayBuffer>}
+   */
+  static async getModel() {
+    const client = SummarizerModel.#getRemoteClient();
+
+    if (!SummarizerModel.#modelRecord) {
+      // Place the records into a promise to prevent any races.
+      SummarizerModel.#modelRecord = (async () => {
+        // Load the wasm binary from remote settings, if it hasn't been already.
+        lazy.console.log(`Getting the summarizer model record.`);
+
+        // TODO - The getMaxVersionRecords should eventually migrated to some kind of
+        // shared utility.
+        const { getMaxVersionRecords } = lazy.TranslationsParent;
+
+        /** @type {WasmRecord[]} */
+        const wasmRecords = await getMaxVersionRecords(client, {
+          // TODO - This record needs to be created with the engine wasm payload.
+          filters: { name: "summarizer-model" },
+          majorVersion: SummarizerModel.MODEL_MAJOR_VERSION,
+        });
+
+        if (wasmRecords.length === 0) {
+          // The remote settings client provides an empty list of records when there is
+          // an error.
+          throw new Error("Unable to get the models from Remote Settings.");
+        }
+
+        if (wasmRecords.length > 1) {
+          SummarizerModel.reportError(
+            new Error("Expected the ml engine to only have 1 record."),
+            wasmRecords
+          );
+        }
+        const [record] = wasmRecords;
+        lazy.console.log(
+          `Using ${record.name}@${record.release} release version ${record.version} first released on Fx${record.fx_release}`,
+          record
+        );
+        return record;
+      })();
+    }
+
+    try {
+      /** @type {{buffer: ArrayBuffer}} */
+      const { buffer } = await client.attachments.download(
+        await SummarizerModel.#modelRecord
+      );
+
+      return buffer;
+    } catch (error) {
+      SummarizerModel.#modelRecord = null;
+      throw error;
+    }
+  }
+
+  /**
+   * Lazily initializes the RemoteSettingsClient.
+   *
+   * @returns {RemoteSettingsClient}
+   */
+  static #getRemoteClient() {
+    if (SummarizerModel.#remoteClient) {
+      return SummarizerModel.#remoteClient;
+    }
+
+    /** @type {RemoteSettingsClient} */
+    const client = lazy.RemoteSettings("ml-model");
+
+    SummarizerModel.#remoteClient = client;
+
+    client.on("sync", async ({ data: { created, updated, deleted } }) => {
+      lazy.console.log(`"sync" event for ml-model`, {
+        created,
+        updated,
+        deleted,
+      });
+
+      // Remove all the deleted records.
+      for (const record of deleted) {
+        await client.attachments.deleteDownloaded(record);
+      }
+
+      // Remove any updated records, and download the new ones.
+      for (const { old: oldRecord } of updated) {
+        await client.attachments.deleteDownloaded(oldRecord);
+      }
+
+      // Do nothing for the created records.
+    });
+
+    return client;
+  }
+}