ml fixes and improvements

apps/weblibre/lib/features/geckoview/features/tabs/domain/repositories/gecko_inference.dart

@@ -202,6 +202,11 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
                   .toList(),
         );
 
+        final idsByTitle = <String, List<String>>{};
+        for (final MapEntry(:key, :value) in unassignedDocumentsInput.entries) {
+          (idsByTitle[value] ??= []).add(key);
+        }
+
         final clusterResult = await clusters.mapNotNull(
           (cluster) => Future.wait(
             cluster.map((clusterTitles) async {
@@ -224,11 +229,7 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
                   },
                 ),
                 tabIds: originalTitles
-                    .map(
+                    .expand((title) => idsByTitle[title] ?? const <String>[])
-                      (title) => unassignedDocumentsInput.entries
-                          .firstWhere((entry) => entry.value == title)
-                          .key,
-                    )
                     .toList(),
               );
             }),
@@ -266,7 +267,8 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
           await _initialLoadComplete.future;
 
           final embeddings = await Result.fromAsync(
-            () async => await _service.generateDocumentEmbeddings(documents),
+            () async =>
+                await _service.generateDocumentEmbeddings(embeddingsToGenerate),
           );
 
           return embeddings;
@@ -276,8 +278,25 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
           return Result.failure(generatedEmbeddings.error!);
         }
 
+        final generated = generatedEmbeddings.value;
+        if (generated.length != embeddingsToGenerate.length) {
+          return Result.failure(
+            ErrorMessage(
+              source: 'Document Embeddings',
+              message: 'Unexpected embedding count',
+              details: {
+                'expected': embeddingsToGenerate.length,
+                'actual': generated.length,
+              },
+            ),
+          );
+        }
+
         for (var i = 0; i < embeddingsToGenerate.length; i++) {
-          embeddings[embeddingsToGenerate[i]] = generatedEmbeddings.value[i];
+          final embedding = generated[i];
+          final document = embeddingsToGenerate[i];
+          embeddings[document] = embedding;
+          _embeddingCache.set(document, embedding);
         }
       }
 
@@ -289,6 +308,15 @@ class GeckoInferenceRepository extends _$GeckoInferenceRepository {
       return Result.failure(
         const ErrorMessage(source: 'Document Embeddings', message: 'Timeout'),
       );
+    } catch (e, s) {
+      return Result.failure(
+        ErrorMessage(
+          source: 'Document Embeddings',
+          message: e.toString(),
+          details: e,
+          stackTrace: s,
+        ),
+      );
     }
   }
 
@@ -456,15 +484,18 @@ Future<List<String>?> containerTabSuggestions(
                 .toList(),
           );
 
-      return suggestedTitles.mapNotNull(
+      final idsByTitle = <String, List<String>>{};
-        (titles) => titles
+      for (final (id, title) in unassignedTitles.value) {
-            .map(
+        (idsByTitle[title] ??= []).add(id);
-              (title) => unassignedTitles.value
+      }
-                  .firstWhere((tab) => tab.$2 == title)
+
-                  .$1,
+      return suggestedTitles.mapNotNull((titles) {
-            )
+        final tabIds = titles
-            .toList(),
+            .expand((title) => idsByTitle[title] ?? const <String>[])
-      );
+            .toList();
+
+        return tabIds.isEmpty ? null : tabIds;
+      });
     }
   }
 

packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/api/ml.js

@@ -4,7 +4,10 @@
 
 "use strict";
 
-const { createEngine } = ChromeUtils.importESModule("chrome://global/content/ml/EngineProcess.sys.mjs");
+const {
+    createEngine,
+    FEATURES,
+} = ChromeUtils.importESModule("chrome://global/content/ml/EngineProcess.sys.mjs");
 
 const ML_TASK_FEATURE_EXTRACTION = "feature-extraction";
 const ML_TASK_TEXT2TEXT = "text2text-generation";
@@ -15,7 +18,8 @@ const SMART_TAB_GROUPING_CONFIG = {
         timeoutMS: 2 * 60 * 1000, // 2 minutes
         taskName: ML_TASK_FEATURE_EXTRACTION,
         featureId: "smart-tab-embedding",
-        backend: "onnx",
+        engineId: FEATURES["smart-tab-embedding"].engineId,
+        backend: "onnx-native",
         fallbackBackend: "onnx",
     },
     topicGeneration: {
@@ -23,7 +27,8 @@ const SMART_TAB_GROUPING_CONFIG = {
         timeoutMS: 2 * 60 * 1000, // 2 minutes
         taskName: ML_TASK_TEXT2TEXT,
         featureId: "smart-tab-topic",
-        backend: "onnx",
+        engineId: FEATURES["smart-tab-topic"].engineId,
+        backend: "onnx-native",
         fallbackBackend: "onnx",
     },
     // dataConfig: {
@@ -117,6 +122,52 @@ function createProgressCallback(modelType, progressEmitter) {
     };
 }
 
+async function createMlEngine(engineConfig, progressCallback) {
+    const {
+        featureId,
+        engineId,
+        dtype,
+        taskName,
+        timeoutMS,
+        modelId,
+        modelRevision,
+        backend,
+        fallbackBackend,
+    } = engineConfig;
+    const initData = {
+        featureId,
+        engineId,
+        dtype,
+        taskName,
+        timeoutMS,
+        modelId,
+        modelRevision,
+        backend,
+    };
+
+    try {
+        return await createEngine(initData, progressCallback);
+    } catch (error) {
+        if (!fallbackBackend || fallbackBackend === backend) {
+            throw error;
+        }
+
+        try {
+            return await createEngine(
+                {
+                    ...initData,
+                    backend: fallbackBackend,
+                },
+                progressCallback
+            );
+        } catch (fallbackError) {
+            throw new Error(
+                `Failed to create ML engine with ${backend} (${error?.message || error}) or ${fallbackBackend} (${fallbackError?.message || fallbackError})`
+            );
+        }
+    }
+}
+
 this.ml = class extends ExtensionAPI {
     constructor(extension) {
         super(extension);
@@ -151,7 +202,7 @@ this.ml = class extends ExtensionAPI {
                         };
 
                         if (isEngineClosed(self.embeddingEngine)) {
-                            self.embeddingEngine = await createEngine(
+                            self.embeddingEngine = await createMlEngine(
                                 SMART_TAB_GROUPING_CONFIG.embedding,
                                 createProgressCallback("Embedding Model", self.progressEmitter)
                             );
@@ -168,30 +219,8 @@ this.ml = class extends ExtensionAPI {
                     },
                     async predictTopic(keywords, documents) {
                         if (isEngineClosed(self.topicEngine)) {
-                            const {
+                            self.topicEngine = await createMlEngine(
-                                featureId,
+                                SMART_TAB_GROUPING_CONFIG.topicGeneration,
-                                engineId,
-                                dtype,
-                                taskName,
-                                timeoutMS,
-                                modelId,
-                                modelRevision,
-                                backend,
-                            } = SMART_TAB_GROUPING_CONFIG.topicGeneration;
-
-                            let initData = {
-                                featureId,
-                                engineId,
-                                dtype,
-                                taskName,
-                                timeoutMS,
-                                modelId,
-                                modelRevision,
-                                backend,
-                            };
-
-                            self.topicEngine = await createEngine(
-                                initData,
                                 createProgressCallback("Topic Generation Model", self.progressEmitter)
                             );
                         }
@@ -221,4 +250,4 @@ this.ml = class extends ExtensionAPI {
             }
         };
     }
-};
+};

packages/flutter_mozilla_components/android/src/main/assets/extensions/ml_engine/background.js

@@ -18,7 +18,7 @@ function sendErrorForRequest(id) {
         port.postMessage({
             "id": id,
             "status": "error",
-            "error": error
+            "error": error?.message || String(error)
         });
     }
 }
@@ -31,27 +31,31 @@ browser.experiments.ml.onProgress.addListener((progressData) => {
 });
 
 port.onMessage.addListener(async (message) => {
-    let requestId = message["id"]
+    const requestId = message["id"];
-    switch (message["action"]) {
-        case "predictDocumentTopic": {
-            const documents = message["args"];
-            const keywords = (documents.length > 1)
-                ? await browser.experiments.nlp.extractKeywords([documents.slice(0, 3).join(" ")])
-                : [[]];
 
-            browser.experiments.ml.predictTopic(keywords[0], documents)
+    try {
-                .then(sendJsonResultForRequest(requestId))
+        switch (message["action"]) {
-                .catch(sendErrorForRequest(requestId));
+            case "predictDocumentTopic": {
+                const documents = message["args"];
+                const keywords = (documents.length > 1)
+                    ? await browser.experiments.nlp.extractKeywords([documents.slice(0, 3).join(" ")])
+                    : [[]];
+                const result = await browser.experiments.ml.predictTopic(keywords[0], documents);
 
-            break;
+                sendJsonResultForRequest(requestId)(result);
-        }
+                break;
-        case "generateDocumentEmbeddings": {
+            }
-            const documents = message["args"];
+            case "generateDocumentEmbeddings": {
-            await browser.experiments.ml.generateEmbeddings(documents)
+                const documents = message["args"];
-                .then(sendJsonResultForRequest(requestId))
+                const result = await browser.experiments.ml.generateEmbeddings(documents);
-                .catch(sendErrorForRequest(requestId));
+
-
+                sendJsonResultForRequest(requestId)(result);
-            break;
+                break;
+            }
+            default:
+                throw new Error(`Unsupported ML action: ${message["action"]}`);
         }
+    } catch (error) {
+        sendErrorForRequest(requestId)(error);
     }
 });

packages/flutter_mozilla_components/android/src/main/kotlin/eu/weblibre/flutter_mozilla_components/feature/MLEngineFeature.kt

@@ -7,7 +7,11 @@
 package eu.weblibre.flutter_mozilla_components.feature
 
 import androidx.annotation.VisibleForTesting
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
@@ -26,10 +30,12 @@ object MLEngineFeature {
     private const val ML_ENGINE_REPORTER_EXTENSION_ID = "ml-engine@weblibre.eu"
     private const val ML_ENGINE_REPORTER_EXTENSION_URL = "resource://android/assets/extensions/ml_engine/"
     private const val ML_ENGINE_REPORTER_MESSAGING_ID = "mlEngine"
+    private const val REQUEST_TIMEOUT_MS = 130_000L
 
     private var nextRequestId: Int = 0
     private val requestHandlers = HashMap<Int, ResultConsumer<JSONObject>>()
     private val mutex = Mutex()
+    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.Default)
 
     // Progress callback for ML operations
     var progressCallback: ((JSONObject) -> Unit)? = null
@@ -44,29 +50,54 @@ object MLEngineFeature {
 
     fun scheduleRequest(command: String, args: Any, callback: ResultConsumer<JSONObject>) {
         val message = JSONObject()
-        message.put("action", command);
+        message.put("action", command)
         message.put("args", args)
 
-        runBlocking {
+        val requestId = runBlocking {
             withContext(Dispatchers.Default) {
                 mutex.withLock {
-                    message.put("id", nextRequestId)
+                    val requestId = nextRequestId
+                    message.put("id", requestId)
 
-                    requestHandlers[nextRequestId] = callback
+                    requestHandlers[requestId] = callback
 
                     nextRequestId += 1
 
-                    extensionController.sendBackgroundMessage(message)
+                    try {
+                        extensionController.sendBackgroundMessage(message)
+                    } catch (throwable: Throwable) {
+                        requestHandlers.remove(requestId)
+                        callback.error(
+                            "ML Engine",
+                            "Failed to schedule request",
+                            throwable.message,
+                        )
+                    }
+
+                    requestId
                 }
             }
         }
+
+        scope.launch {
+            delay(REQUEST_TIMEOUT_MS)
+            val handler = mutex.withLock {
+                requestHandlers.remove(requestId)
+            }
+
+            handler?.error(
+                "ML Engine",
+                "Request timed out",
+                "No response received for $command",
+            )
+        }
     }
 
     private class PrefManagerReporterBackgroundMessageHandler() : MessageHandler {
         override fun onPortMessage(message: Any, port: Port) {
             runBlocking {
                 withContext(Dispatchers.Default) {
-                    val messageJSON = message as JSONObject;
+                    val messageJSON = message as JSONObject
 
                     // Check if this is a progress message
                     if (messageJSON.has("type") && messageJSON.getString("type") == "mlProgress") {
@@ -86,7 +117,7 @@ object MLEngineFeature {
                             handler?.error(
                                 "ML Engine",
                                 "Failed to perform operation",
-                                message.getString("error")
+                                messageJSON.optString("error", "Unknown ML engine error")
                             )
                         }
                     }