[Function] Add heartbeat to service worker (#400)

## Overview
This PR adds heartbeat event in web service worker so that the client
can monitor its status and respond correspondingly.

## Primary Changes
- Update `{ type: "keepAlive" }` event to `{ kind: "keepAlive" }` to
keep all event format consistent
- Add heartbeat event in web service worker handler to report back its
status

## Testing
<img width="616" alt="Screenshot 2024-05-15 at 2 01 51 AM"
src="https://github.com/mlc-ai/web-llm/assets/23090573/42fcd7b0-f8f2-4b23-80c3-664426853161">

Author: Neet-Nestor

src/message.ts

@@ -16,7 +16,8 @@ type RequestKind = (
   "interruptGenerate" | "unload" | "resetChat" | "init" |
   "initProgressCallback" | "generateProgressCallback" | "getMaxStorageBufferBindingSize" |
   "getGPUVendor" | "forwardTokensAndSample" | "chatCompletionNonStreaming" | "getMessage" |
-  "chatCompletionStreamInit" | "chatCompletionStreamNextChunk" | "customRequest");
+  "chatCompletionStreamInit" | "chatCompletionStreamNextChunk" | "customRequest" | 'keepAlive' | 'heartbeat');
+
 export interface ReloadParams {
   modelId: string;
   chatOpts?: ChatOptions;

src/service_worker.ts

@@ -139,7 +139,7 @@ export async function CreateServiceWorkerEngine(
   engineConfig?: EngineConfig,
   keepAliveMs: number = 10000
 ): Promise<ServiceWorkerEngine> {
-  const serviceWorkerEngine = new ServiceWorkerEngine();
+  const serviceWorkerEngine = new ServiceWorkerEngine(keepAliveMs);
   serviceWorkerEngine.setInitProgressCallback(
     engineConfig?.initProgressCallback
   );
@@ -148,9 +148,6 @@ export async function CreateServiceWorkerEngine(
     engineConfig?.chatOpts,
     engineConfig?.appConfig
   );
-  setInterval(() => {
-    serviceWorkerEngine.keepAlive();
-  }, keepAliveMs);
   return serviceWorkerEngine;
 }
 
@@ -191,15 +188,18 @@ class PortAdapter implements ChatWorker {
 export class ServiceWorkerEngine extends WebWorkerEngine {
   port: chrome.runtime.Port;
 
-  constructor() {
+  constructor(keepAliveMs: number = 10000) {
     let port = chrome.runtime.connect({ name: "web_llm_service_worker" });
     let chatWorker = new PortAdapter(port);
     super(chatWorker);
     this.port = port;
+    setInterval(() => {
+      this.keepAlive();
+    }, keepAliveMs);
   }
 
   keepAlive() {
-    this.port.postMessage({ type: "keepAlive" });
+    this.worker.postMessage({ kind: "keepAlive" });
   }
 
   /**

src/web_service_worker.ts

@@ -74,6 +74,16 @@ export class ServiceWorkerEngineHandler extends EngineWorkerHandler {
     const msgEvent = event as MessageEvent;
     const msg = msgEvent.data as WorkerMessage;
 
+    if (msg.kind === "keepAlive") {
+      const msg: WorkerMessage = {
+        kind: "heartbeat",
+        uuid: msgEvent.data.uuid,
+        content: "",
+      };
+      this.postMessageInternal(msg);
+      return;
+    }
+
     if (msg.kind === "init") {
       this.handleTask(msg.uuid, async () => {
         const params = msg.content as ReloadParams;
@@ -114,7 +124,7 @@ export class ServiceWorkerEngineHandler extends EngineWorkerHandler {
       });
       return;
     }
-    super.onmessage(event);
+    super.onmessage(msg);
   }
 }
 
@@ -147,13 +157,36 @@ export async function CreateServiceWorkerEngine(
  * A client of Engine that exposes the same interface
  */
 export class ServiceWorkerEngine extends WebWorkerEngine {
-  constructor(worker: ChatWorker) {
+  missedHeatbeat = 0
+
+  constructor(worker: ChatWorker, keepAliveMs=10000) {
     super(worker);
-    clientBroadcastChannel.onmessage = this.onmessage.bind(this);
+    clientBroadcastChannel.onmessage = (event) => {
+      try {
+        this.onevent.bind(this)(event)
+      } catch (err: any) {
+        // This is expected to throw if user has multiple windows open
+        if (!err.message.startsWith("return from a unknown uuid")) {
+          console.error("CreateWebServiceWorkerEngine.onmessage", err);
+        }
+      }
+    };
+    setInterval(() => {
+      this.keepAlive();
+    }, keepAliveMs);
   }
 
   keepAlive() {
-    this.worker.postMessage({ type: "keepAlive" });
+    this.worker.postMessage({ kind: "keepAlive" });
+  }
+
+  onevent(event: MessageEvent): void {
+    const msg = event.data as WorkerMessage;
+    if (msg.kind === "heartbeat") {
+      this.missedHeatbeat = 0
+      return;
+    }
+    this.onmessage(msg);
   }
 
   /**