1 files changed, 185 insertions, 3 deletions

diff --git a/toolkit/components/ml/content/EngineProcess.sys.mjs b/toolkit/components/ml/content/EngineProcess.sys.mjs
index 36a9381192..0fe6403cc8 100644
--- a/toolkit/components/ml/content/EngineProcess.sys.mjs
+++ b/toolkit/components/ml/content/EngineProcess.sys.mjs
@@ -2,10 +2,17 @@
  * 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/. */
 
+// known to be loaded early in the startup process, and should be loaded eagerly
+import { AppConstants } from "resource://gre/modules/AppConstants.sys.mjs";
+
 const lazy = {};
-ChromeUtils.defineESModuleGetters(lazy, {
-  HiddenFrame: "resource://gre/modules/HiddenFrame.sys.mjs",
-});
+ChromeUtils.defineESModuleGetters(
+  lazy,
+  {
+    HiddenFrame: "resource://gre/modules/HiddenFrame.sys.mjs",
+  },
+  { global: "current" }
+);
 
 /**
  * @typedef {import("../actors/MLEngineParent.sys.mjs").MLEngineParent} MLEngineParent
@@ -16,6 +23,172 @@ ChromeUtils.defineESModuleGetters(lazy, {
  */
 
 /**
+ * This class encapsulates the options for a pipeline process.
+ */
+export class PipelineOptions {
+  /**
+   * The name of the task the pipeline is configured for.
+   *
+   * @type {?string}
+   */
+  taskName = null;
+
+  /**
+   * The maximum amount of time in milliseconds the pipeline should wait for a response.
+   *
+   * @type {?number}
+   */
+  timeoutMS = null;
+
+  /**
+   * The root URL of the model hub where models are hosted.
+   *
+   * @type {?string}
+   */
+  modelHubRootUrl = null;
+
+  /**
+   * A template URL for building the full URL for the model.
+   *
+   * @type {?string}
+   */
+  modelHubUrlTemplate = null;
+
+  /**
+   * The identifier for the specific model to be used by the pipeline.
+   *
+   * @type {?string}
+   */
+  modelId = null;
+
+  /**
+   * The revision for the specific model to be used by the pipeline.
+   *
+   * @type {?string}
+   */
+  modelRevision = null;
+
+  /**
+   * The identifier for the tokenizer associated with the model, used for pre-processing inputs.
+   *
+   * @type {?string}
+   */
+  tokenizerId = null;
+
+  /**
+   * The revision for the tokenizer associated with the model, used for pre-processing inputs.
+   *
+   * @type {?string}
+   */
+  tokenizerRevision = null;
+
+  /**
+   * The identifier for any processor required by the model, used for additional input processing.
+   *
+   * @type {?string}
+   */
+  processorId = null;
+
+  /**
+   * The revision for any processor required by the model, used for additional input processing.
+   *
+   * @type {?string}
+   */
+
+  processorRevision = null;
+
+  /**
+   * The log level used in the worker
+   *
+   * @type {?string}
+   */
+  logLevel = null;
+
+  /**
+   * Name of the runtime wasm file
+   *
+   * @type {?string}
+   */
+  runtimeFilename = null;
+
+  /**
+   * Create a PipelineOptions instance.
+   *
+   * @param {object} options - The options for the pipeline. Must include mandatory fields.
+   */
+  constructor(options) {
+    this.updateOptions(options);
+  }
+
+  /**
+   * Updates multiple options at once.
+   *
+   * @param {object} options - An object containing the options to update.
+   * @throws {Error} Throws an error if an invalid option is provided.
+   */
+  updateOptions(options) {
+    const allowedKeys = [
+      "taskName",
+      "modelHubRootUrl",
+      "modelHubUrlTemplate",
+      "timeoutMS",
+      "modelId",
+      "modelRevision",
+      "tokenizerId",
+      "tokenizerRevision",
+      "processorId",
+      "processorRevision",
+      "logLevel",
+      "runtimeFilename",
+    ];
+
+    Object.keys(options).forEach(key => {
+      if (allowedKeys.includes(key)) {
+        this[key] = options[key]; // Use bracket notation to access setter
+      } else {
+        throw new Error(`Invalid option: ${key}`);
+      }
+    });
+  }
+
+  /**
+   * Returns an object containing all current options.
+
+   * @returns {object} An object with the current options.
+   */
+  getOptions() {
+    return {
+      taskName: this.taskName,
+      modelHubRootUrl: this.modelHubRootUrl,
+      modelHubUrlTemplate: this.modelHubUrlTemplate,
+      timeoutMS: this.timeoutMS,
+      modelId: this.modelId,
+      modelRevision: this.modelRevision,
+      tokenizerId: this.tokenizerId,
+      tokenizerRevision: this.tokenizerRevision,
+      processorId: this.processorId,
+      processorRevision: this.processorRevision,
+      logLevel: this.logLevel,
+      runtimeFilename: this.runtimeFilename,
+    };
+  }
+
+  /**
+   * Updates the given configuration object with the options.
+   *
+   * @param {object} config - The configuration object to be updated.
+   */
+  applyToConfig(config) {
+    const options = this.getOptions();
+    Object.keys(options).forEach(key => {
+      if (options[key] !== null) {
+        config[key] = options[key];
+      }
+    });
+  }
+}
+
+/**
  * This class controls the life cycle of the engine process used both in the
  * Translations engine and the MLEngine component.
  */
@@ -68,6 +241,15 @@ export class EngineProcess {
    * @returns {Promise<MLEngineParent>}
    */
   static async getMLEngineParent() {
+    // Bug 1890946 - enable the inference engine in release
+    if (!AppConstants.NIGHTLY_BUILD) {
+      throw new Error("MLEngine is only available in Nightly builds.");
+    }
+    // the pref is off by default
+    if (!Services.prefs.getBoolPref("browser.ml.enable")) {
+      throw new Error("MLEngine is disabled. Check the browser.ml prefs.");
+    }
+
     if (!this.mlEngineParent) {
       this.mlEngineParent = this.#attachBrowser({
         id: "ml-engine-browser",