feat: fall back to Web Worker if WebGPU is not available in Service W…

…orker

app/client/webllm.ts

@@ -34,17 +34,18 @@ export class WebLLMApi implements LLMApi {
   private initialized = false;
   webllm: WebLLMHandler;
 
-  constructor(logLevel: LogLevel = "WARN") {
-    if ("serviceWorker" in navigator && navigator.serviceWorker.controller) {
-      log.info("Service Worker API is available and in use.");
+  constructor(
+    type: "serviceWorker" | "webWorker",
+    logLevel: LogLevel = "WARN",
+  ) {
+    if (type === "serviceWorker") {
+      log.info("Create ServiceWorkerMLCEngine");
       this.webllm = {
         type: "serviceWorker",
         engine: new ServiceWorkerMLCEngine(KEEP_ALIVE_INTERVAL),
       };
     } else {
-      log.info(
-        "Service Worker API is unavailable. Falling back to use web worker.",
-      );
+      log.info("Create WebWorkerMLCEngine");
       this.webllm = {
         type: "webWorker",
         engine: new WebWorkerMLCEngine(

app/components/home.tsx

@@ -167,36 +167,60 @@ const useWebLLM = () => {
   // Initialize WebLLM engine
   useEffect(() => {
     if ("serviceWorker" in navigator) {
+      log.info("Service Worker API is available and in use.");
       navigator.serviceWorker.ready.then(() => {
-        setWebLLM(new WebLLMApi(config.logLevel));
-
-        // Verify service worker has been activated
-        const heartbeatCallback = (event: MessageEvent) => {
-          const msg = event.data;
-          if (msg.kind === "heartbeat") {
-            log.info("Confirmed messages from Service Worker. Starting app...");
-            setWebllmAlive(true);
+        log.info("Service Worker is activated.");
+        // Check whether WebGPU is available in Service Worker
+        const request = {
+          kind: "checkWebGPUAvilability",
+          uuid: crypto.randomUUID(),
+          content: "",
+        };
+        const webGPUCheckCallback = (event: MessageEvent) => {
+          const message = event.data;
+          if (message.kind === "return" && message.uuid === request.uuid) {
+            const isWebGPUAvailable = message.content;
+            log.info(
+              isWebGPUAvailable
+                ? "Service Worker has WebGPU Available."
+                : "Service Worker does not have available WebGPU.",
+            );
+            if (!webllm && !isWebllmActive) {
+              setWebLLM(
+                new WebLLMApi(
+                  isWebGPUAvailable ? "serviceWorker" : "webWorker",
+                  config.logLevel,
+                ),
+              );
+              setWebllmAlive(true);
+            }
             navigator.serviceWorker.removeEventListener(
               "message",
-              heartbeatCallback,
+              webGPUCheckCallback,
             );
           }
         };
-        navigator.serviceWorker.addEventListener("message", heartbeatCallback);
-        navigator.serviceWorker.controller?.postMessage({
-          kind: "keepAlive",
-          uuid: crypto.randomUUID(),
-        });
+        navigator.serviceWorker.addEventListener(
+          "message",
+          webGPUCheckCallback,
+        );
+        navigator.serviceWorker.controller?.postMessage(request);
       });
     } else {
-      setWebLLM(new WebLLMApi(config.logLevel));
+      log.info(
+        "Service Worker API is unavailable. Falling back to use web worker.",
+      );
+      setWebLLM(new WebLLMApi("webWorker", config.logLevel));
       setWebllmAlive(true);
     }
 
     // If service worker registration timeout
     setTimeout(() => {
       if (!navigator.serviceWorker?.controller && !isWebllmActive && !webllm) {
-        setWebLLM(new WebLLMApi(config.logLevel));
+        log.info(
+          "Service Worker activation is timed out. Falling back to use web worker.",
+        );
+        setWebLLM(new WebLLMApi("webWorker", config.logLevel));
         setWebllmAlive(true);
       }
     }, 5_000);

app/worker/service-worker.ts

@@ -2,17 +2,46 @@ import { ServiceWorkerMLCEngineHandler, MLCEngine } from "@neet-nestor/web-llm";
 import { defaultCache } from "@serwist/next/worker";
 import type { PrecacheEntry, SerwistGlobalConfig } from "serwist";
 import { CacheFirst, ExpirationPlugin, Serwist } from "serwist";
-import log from "loglevel";
 
 declare const self: ServiceWorkerGlobalScope;
 const CHATGPT_NEXT_WEB_CACHE = "chatgpt-next-web-cache";
 const engine = new MLCEngine();
 let handler: ServiceWorkerMLCEngineHandler;
 
+async function checkGPUAvailablity() {
+  if (!("gpu" in navigator)) {
+    console.log("Service Worker: Web-LLM Engine Activated");
+    return false;
+  }
+  const adapter = await navigator.gpu.requestAdapter();
+  if (!adapter) {
+    console.log("Service Worker: Web-LLM Engine Activated");
+    return false;
+  }
+  return true;
+}
+
 self.addEventListener("message", (event) => {
   if (!handler) {
     handler = new ServiceWorkerMLCEngineHandler(engine);
-    log.info("Service Worker: Web-LLM Engine Activated");
+    console.log("Service Worker: Web-LLM Engine Activated");
+  }
+
+  const msg = event.data;
+  if (msg.kind === "checkWebGPUAvilability") {
+    console.log("Service Worker: Web-LLM Engine Activated");
+    checkGPUAvailablity().then((gpuAvailable) => {
+      console.log(
+        "Service Worker: WebGPU is " +
+          (gpuAvailable ? "available" : "unavailable"),
+      );
+      const reply = {
+        kind: "return",
+        uuid: msg.uuid,
+        content: gpuAvailable,
+      };
+      event.source?.postMessage(reply);
+    });
   }
 });
 
@@ -30,7 +59,7 @@ self.addEventListener("install", (event) => {
 self.addEventListener("activate", (event) => {
   if (!handler) {
     handler = new ServiceWorkerMLCEngineHandler(engine);
-    log.info("Service Worker: Web-LLM Engine Activated");
+    console.log("Service Worker: Web-LLM Engine Activated");
   }
 });
 

package.json

@@ -49,6 +49,7 @@
     "@types/react-dom": "^18.2.7",
     "@types/react-katex": "^3.0.0",
     "@types/spark-md5": "^3.0.4",
+    "@webgpu/types": "^0.1.42",
     "cross-env": "^7.0.3",
     "eslint": "^8.49.0",
     "eslint-config-next": "13.4.19",

tsconfig.json

@@ -1,7 +1,7 @@
 {
   "compilerOptions": {
     "target": "ES2015",
-    "types": ["@serwist/next/typings"],
+    "types": ["@serwist/next/typings", "@webgpu/types"],
     "lib": ["dom", "dom.iterable", "esnext", "webworker"],
     "allowJs": true,
     "skipLibCheck": true,

yarn.lock

@@ -2889,6 +2889,11 @@
     "@webassemblyjs/ast" "1.11.6"
     "@xtuc/long" "4.2.2"
 
+"@webgpu/types@^0.1.42":
+  version "0.1.42"
+  resolved "https://registry.yarnpkg.com/@webgpu/types/-/types-0.1.42.tgz#02ce82ae7d486521064eda2e9c7e38697136afe9"
+  integrity sha512-uvJtt4OD1Vjdebrrz3kNLgpOicYbikwnM8WPG6YD2lkCOHDtPdEtCINJFIFtbOCtPfA8SreR/vKyUNbAt92IwQ==
+
 "@xtuc/ieee754@^1.2.0":
   version "1.2.0"
   resolved "https://registry.npmmirror.com/@xtuc/ieee754/-/ieee754-1.2.0.tgz#eef014a3145ae477a1cbc00cd1e552336dceb790"