feat: Replace BroadcastChannel API with Client API in service worker

app/client/webllm.ts

@@ -4,23 +4,27 @@ import {
   prebuiltAppConfig,
   ChatCompletionMessageParam,
   ServiceWorkerEngine,
-  ServiceWorker,
   ChatCompletionChunk,
   ChatCompletion,
 } from "@neet-nestor/web-llm";
 
 import { ChatOptions, LLMApi, LLMConfig, RequestMessage } from "./api";
 
-const KEEP_ALIVE_INTERVAL = 10000;
+const KEEP_ALIVE_INTERVAL = 10_000;
 
 export class WebLLMApi implements LLMApi {
   private llmConfig?: LLMConfig;
-  engine?: ServiceWorkerEngine;
+  engine: ServiceWorkerEngine;
 
   constructor() {
-    this.engine = new ServiceWorkerEngine(new ServiceWorker());
-    this.engine.keepAlive(
-      window.location.href + "ping.txt",
+    if (!("serviceWorker" in navigator)) {
+      throw Error("Service worker API is not available");
+    }
+    if (!navigator.serviceWorker.controller) {
+      throw Error("There is no active service worker");
+    }
+    this.engine = new ServiceWorkerEngine(
+      navigator.serviceWorker.controller,
       KEEP_ALIVE_INTERVAL,
     );
   }
@@ -29,44 +33,13 @@ export class WebLLMApi implements LLMApi {
     if (!this.llmConfig) {
       throw Error("llmConfig is undefined");
     }
-    if (!this.engine) {
-      this.engine = new ServiceWorkerEngine(new ServiceWorker());
-    }
-    let hasResponse = false;
     this.engine.setInitProgressCallback((report: InitProgressReport) => {
       onUpdate?.(report.text, report.text);
-      hasResponse = true;
     });
-    let initRequest = this.engine.init(this.llmConfig.model, this.llmConfig, {
+    await this.engine.init(this.llmConfig.model, this.llmConfig, {
       ...prebuiltAppConfig,
       useIndexedDBCache: this.llmConfig.cache === "index_db",
     });
-    // In case the service worker is dead, init will halt indefinitely
-    // so we manually retry if timeout
-    let retry = 0;
-    let engine = this.engine;
-    let llmConfig = this.llmConfig;
-    let retryInterval: NodeJS.Timeout;
-
-    await new Promise<void>((resolve, reject) => {
-      retryInterval = setInterval(() => {
-        if (hasResponse) {
-          clearInterval(retryInterval);
-          initRequest.then(resolve);
-          return;
-        }
-        if (retry >= 5) {
-          clearInterval(retryInterval);
-          reject("Model initialization timed out for too many times");
-          return;
-        }
-        retry += 1;
-        initRequest = engine.init(llmConfig.model, llmConfig, {
-          ...prebuiltAppConfig,
-          useIndexedDBCache: llmConfig.cache === "index_db",
-        });
-      }, 5000);
-    });
   }
 
   isConfigChanged(config: LLMConfig) {
@@ -81,15 +54,12 @@ export class WebLLMApi implements LLMApi {
   }
 
   async chat(options: ChatOptions): Promise<void> {
-    // in case the service worker is dead, revive it by firing a fetch event
-    fetch("/ping.txt");
-
     if (this.isDifferentConfig(options.config)) {
       this.llmConfig = { ...(this.llmConfig || {}), ...options.config };
       try {
         await this.initModel(options.onUpdate);
       } catch (e) {
-        console.error("Error in initModel", e);
+        console.error("Error while initializing the model", e);
       }
     }
 

app/service-worker.ts

@@ -64,8 +64,7 @@ self.addEventListener("activate", (event) => {
   }
 });
 
-self.addEventListener("fetch", (event) => {
-  console.log("sw fetch handler");
+self.addEventListener("message", (event) => {
   if (!handler) {
     handler = new ServiceWorkerEngineHandler(engine);
     console.log("Service Worker: Web-LLM Engine Activated");

package.json

@@ -17,7 +17,7 @@
   "dependencies": {
     "@fortaine/fetch-event-source": "^3.0.6",
     "@hello-pangea/dnd": "^16.5.0",
-    "@neet-nestor/web-llm": "0.2.41",
+    "@neet-nestor/web-llm": "^0.2.44",
     "@next/third-parties": "^14.1.0",
     "@serwist/next": "^9.0.2",
     "@svgr/webpack": "^6.5.1",

yarn.lock

@@ -2089,10 +2089,10 @@
     "@jridgewell/resolve-uri" "^3.1.0"
     "@jridgewell/sourcemap-codec" "^1.4.14"
 
-"@neet-nestor/[email protected].41":
-  version "0.2.41"
-  resolved "https://registry.yarnpkg.com/@neet-nestor/web-llm/-/web-llm-0.2.41.tgz#ff90892d0dc20165153eebd5ea4e92d7e7308fe8"
-  integrity sha512-AEmo4I8PoKGfn486oo//18YZIw6VUdpAKlkSoRMY0EbCvdur8gZ3Y6mi1+YKwHk6FJk6rmpOStn8KC4tmjwU0Q==
+"@neet-nestor/web-llm@^0.2.44":
+  version "0.2.44"
+  resolved "https://registry.yarnpkg.com/@neet-nestor/web-llm/-/web-llm-0.2.44.tgz#1e6ca4d578cb9ceefccd795bcf7e9140125e8a18"
+  integrity sha512-fIvGelFCGtzVb4dN68zKomtVjKIbssGOXGM91Gry9HyYfUjRdq5cxBd1Xg2d1yYtce4twesowuy1fqGX6VBZOg==
 
 "@next/[email protected]":
   version "13.4.9"