[Feature] Replace BroadcastChannel API with Client API in Service Worker (#401)

## Overview

There are many different APIs available for the two-way communication
between service worker and web pages. According to the suggestion of
web.dev, previously we used the simple BroadcastChannel API.
https://web.dev/articles/two-way-communication-guide

However, if the service worker has been killed by the browser, messages
sent via the BroadcastChannel API cannot revive it. This causes unstable
service worker life while users are using the webapp.

This PR replaces BroadcastChannel API with the more primitive Client API
as we found that messages sent via Client API will revive a stopped
service worker. This ensures the normal functioning of our keep-alive
mechanism.

## Primary Change

- service_worker.ts:
- Replace `BroadcastChannel` with Client API
(`navigator.serviceWorker.controller.postMessage()` and
`client.postMessage()`)
- Add `clientRegistry` to remember mapping between incoming messages and
client ids
  - Add 
- Rename files (the export names are kept the same):
  - `web_service_worker.ts` -> `service_worker.ts`
  - `service_worker.ts` -> `extension_service_worker.ts`

## Testing

- https://chat.webllm.ai/
- `examples/service-worker`

The chat webapp is able to correctly keeping service worker alive after
this change.

Author: Neet-Nestor

examples/service-worker/package.json

@@ -3,7 +3,7 @@
     "version": "0.1.0",
     "private": true,
     "scripts": {
-        "start": "parcel src/index.html  --port 3000",
+        "start": "parcel src/index.html --port 3000",
         "build": "parcel build src/index.html --dist-dir lib"
     },
     "devDependencies": {
@@ -15,6 +15,6 @@
         "url": "^0.11.3"
     },
     "dependencies": {
-        "@mlc-ai/web-llm": "file:../../lib"
+        "@mlc-ai/web-llm": "file:../../lib/"
     }
 }

examples/service-worker/src/main.ts

@@ -77,7 +77,7 @@ async function mainStreaming() {
   };
   const selectedModel = "Llama-3-8B-Instruct-q4f32_1";
 
-  const engine: webllm.EngineInterface =
+  const engine: webllm.WebServiceWorkerEngine =
     await webllm.CreateWebServiceWorkerEngine(selectedModel, {
       initProgressCallback: initProgressCallback,
     });

package-lock.json

@@ -30,6 +30,7 @@
     "@rollup/plugin-node-resolve": "^13.0.4",
     "@types/chrome": "^0.0.266",
     "@types/jest": "^29.5.11",
+    "@types/serviceworker": "^0.0.86",
     "@typescript-eslint/eslint-plugin": "^5.59.6",
     "@typescript-eslint/parser": "^5.59.6",
     "@webgpu/types": "^0.1.24",
@@ -45,4 +46,4 @@
     "tvmjs": "file:./tvm_home/web",
     "typescript": "^4.9.5"
   }
-}
+}

package.json

@@ -1,46 +1,37 @@
 import * as tvmjs from "tvmjs";
 import { AppConfig, ChatOptions, EngineConfig } from "./config";
-import { ReloadParams, WorkerMessage } from "./message";
+import { ReloadParams, WorkerRequest } from "./message";
 import { EngineInterface } from "./types";
 import {
+  ChatWorker,
   EngineWorkerHandler,
   WebWorkerEngine,
   PostMessageHandler,
-  ChatWorker,
 } from "./web_worker";
 import { areAppConfigsEqual, areChatOptionsEqual } from "./utils";
 
-const BROADCAST_CHANNEL_SERVICE_WORKER_ID = "@mlc-ai/web-llm-sw";
-const BROADCAST_CHANNEL_CLIENT_ID = "@mlc-ai/web-llm-client";
-export const serviceWorkerBroadcastChannel = new BroadcastChannel(
-  BROADCAST_CHANNEL_SERVICE_WORKER_ID
-);
-export const clientBroadcastChannel = new BroadcastChannel(
-  BROADCAST_CHANNEL_CLIENT_ID
-);
-
 /**
- * PostMessageHandler wrapper for sending message from service worker back to client
+ * A post message handler that sends messages to a chrome.runtime.Port.
  */
-class ClientPostMessageHandler implements PostMessageHandler {
-  postMessage(message: any) {
-    clientBroadcastChannel.postMessage(message);
-  }
-}
+export class PortPostMessageHandler implements PostMessageHandler {
+  port: chrome.runtime.Port;
+  enabled: boolean = true;
 
-/**
- * PostMessageHandler wrapper for sending message from client to service worker
- */
-class ServiceWorker implements ChatWorker {
-  constructor() {
-    serviceWorkerBroadcastChannel.onmessage = this.onmessage;
+  constructor(port: chrome.runtime.Port) {
+    this.port = port;
   }
 
-  // ServiceWorkerEngine will later overwrite this
-  onmessage() {}
+  /**
+   * Close the PortPostMessageHandler. This will prevent any further messages
+   */
+  close() {
+    this.enabled = false;
+  }
 
-  postMessage(message: any) {
-    serviceWorkerBroadcastChannel.postMessage(message);
+  postMessage(event: any) {
+    if (this.enabled) {
+      this.port.postMessage(event);
+    }
   }
 }
 
@@ -65,25 +56,29 @@ export class ServiceWorkerEngineHandler extends EngineWorkerHandler {
   chatOpts?: ChatOptions;
   appConfig?: AppConfig;
 
-  constructor(engine: EngineInterface) {
-    super(engine, new ClientPostMessageHandler());
-    serviceWorkerBroadcastChannel.onmessage = this.onmessage.bind(this);
+  constructor(engine: EngineInterface, port: chrome.runtime.Port) {
+    let portHandler = new PortPostMessageHandler(port);
+    super(engine, portHandler);
+
+    port.onDisconnect.addListener(() => {
+      portHandler.close();
+    });
+  }
+
+  setPort(port: chrome.runtime.Port) {
+    let portHandler = new PortPostMessageHandler(port);
+    this.setPostMessageHandler(portHandler);
+    port.onDisconnect.addListener(() => {
+      portHandler.close();
+    });
   }
 
   onmessage(event: any): void {
-    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);
+    if (event.type === "keepAlive") {
       return;
     }
 
+    const msg = event as WorkerRequest;
     if (msg.kind === "init") {
       this.handleTask(msg.uuid, async () => {
         const params = msg.content as ReloadParams;
@@ -124,7 +119,7 @@ export class ServiceWorkerEngineHandler extends EngineWorkerHandler {
       });
       return;
     }
-    super.onmessage(msg);
+    super.onmessage(event);
   }
 }
 
@@ -134,14 +129,17 @@ export class ServiceWorkerEngineHandler extends EngineWorkerHandler {
  * @param modelId The model to load, needs to either be in `webllm.prebuiltAppConfig`, or in
  * `engineConfig.appConfig`.
  * @param engineConfig Optionally configures the engine, see `webllm.EngineConfig` for more.
+ * @param keepAliveMs The interval to send keep alive messages to the service worker.
+ * See [Service worker lifecycle](https://developer.chrome.com/docs/extensions/develop/concepts/service-workers/lifecycle#idle-shutdown)
+ * The default is 10s.
  * @returns An initialized `WebLLM.ServiceWorkerEngine` with `modelId` loaded.
  */
 export async function CreateServiceWorkerEngine(
   modelId: string,
-  engineConfig?: EngineConfig
+  engineConfig?: EngineConfig,
+  keepAliveMs: number = 10000
 ): Promise<ServiceWorkerEngine> {
-  await navigator.serviceWorker.ready;
-  const serviceWorkerEngine = new ServiceWorkerEngine(new ServiceWorker());
+  const serviceWorkerEngine = new ServiceWorkerEngine(keepAliveMs);
   serviceWorkerEngine.setInitProgressCallback(
     engineConfig?.initProgressCallback
   );
@@ -153,24 +151,48 @@ export async function CreateServiceWorkerEngine(
   return serviceWorkerEngine;
 }
 
+class PortAdapter implements ChatWorker {
+  port: chrome.runtime.Port;
+  private _onmessage!: (message: any) => void;
+
+  constructor(port: chrome.runtime.Port) {
+    this.port = port;
+    this.port.onMessage.addListener(this.handleMessage.bind(this));
+  }
+
+  // Wrapper to handle incoming messages and delegate to onmessage if available
+  private handleMessage(message: any) {
+    if (this._onmessage) {
+      this._onmessage(message);
+    }
+  }
+
+  // Getter and setter for onmessage to manage adding/removing listeners
+  get onmessage(): (message: any) => void {
+    return this._onmessage;
+  }
+
+  set onmessage(listener: (message: any) => void) {
+    this._onmessage = listener;
+  }
+
+  // Wrap port.postMessage to maintain 'this' context
+  postMessage = (message: any): void => {
+    this.port.postMessage(message);
+  };
+}
+
 /**
  * A client of Engine that exposes the same interface
  */
 export class ServiceWorkerEngine extends WebWorkerEngine {
-  missedHeatbeat = 0
-
-  constructor(worker: ChatWorker, keepAliveMs=10000) {
-    super(worker);
-    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);
-        }
-      }
-    };
+  port: chrome.runtime.Port;
+
+  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);
@@ -180,15 +202,6 @@ export class ServiceWorkerEngine extends WebWorkerEngine {
     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);
-  }
-
   /**
    * Initialize the chat with a model.
    *
@@ -205,7 +218,7 @@ export class ServiceWorkerEngine extends WebWorkerEngine {
     chatOpts?: ChatOptions,
     appConfig?: AppConfig
   ): Promise<void> {
-    const msg: WorkerMessage = {
+    const msg: WorkerRequest = {
       kind: "init",
       uuid: crypto.randomUUID(),
       content: {

src/web_service_worker.ts renamed to src/extension_service_worker.ts

@@ -33,24 +33,23 @@ export {
 } from "./web_worker";
 
 export {
-  WorkerMessage,
+  WorkerRequest,
+  WorkerResponse,
   CustomRequestParams
 } from "./message"
 
-// TODO: Rename to ExtensionServiceWorker 
-export {
-  ServiceWorkerEngineHandler,
-  ServiceWorkerEngine,
-  CreateServiceWorkerEngine,
-} from './service_worker'
-
-// TODO: Rename to ServiceWorker 
+// TODO: Rename classes to ServiceWorker 
 export {
   ServiceWorkerEngineHandler as WebServiceWorkerEngineHandler,
   ServiceWorkerEngine as WebServiceWorkerEngine,
   CreateServiceWorkerEngine as CreateWebServiceWorkerEngine,
-  serviceWorkerBroadcastChannel,
-  clientBroadcastChannel,
-} from "./web_service_worker";
+} from "./service_worker";
+
+// TODO: Rename classes to ExtensionServiceWorker 
+export {
+  ServiceWorkerEngineHandler,
+  ServiceWorkerEngine,
+  CreateServiceWorkerEngine,
+} from './extension_service_worker'
 
 export * from './openai_api_protocols/index';

src/index.ts

@@ -4,19 +4,36 @@ import {
   ChatCompletionRequestStreaming,
   ChatCompletionRequestNonStreaming,
   ChatCompletion,
-  ChatCompletionChunk
+  ChatCompletionChunk,
 } from "./openai_api_protocols/index";
 
 /**
  * Message kind used by worker
  */
-type RequestKind = (
-  "return" | "throw" |
-  "reload" | "generate" | "runtimeStatsText" |
-  "interruptGenerate" | "unload" | "resetChat" | "init" |
-  "initProgressCallback" | "generateProgressCallback" | "getMaxStorageBufferBindingSize" |
-  "getGPUVendor" | "forwardTokensAndSample" | "chatCompletionNonStreaming" | "getMessage" |
-  "chatCompletionStreamInit" | "chatCompletionStreamNextChunk" | "customRequest" | 'keepAlive' | 'heartbeat');
+type RequestKind =
+  | "reload"
+  | "generate"
+  | "runtimeStatsText"
+  | "interruptGenerate"
+  | "unload"
+  | "resetChat"
+  | "init"
+  | "getMaxStorageBufferBindingSize"
+  | "getGPUVendor"
+  | "forwardTokensAndSample"
+  | "chatCompletionNonStreaming"
+  | "getMessage"
+  | "chatCompletionStreamInit"
+  | "chatCompletionStreamNextChunk"
+  | "customRequest"
+  | "keepAlive"
+  | "heartbeat";
+
+type ResponseKind =
+  | "return"
+  | "throw"
+  | "initProgressCallback"
+  | "generateProgressCallback";
 
 export interface ReloadParams {
   modelId: string;
@@ -51,28 +68,54 @@ export interface CustomRequestParams {
   requestMessage: string;
 }
 export type MessageContent =
-  GenerateProgressCallbackParams |
-  ReloadParams |
-  GenerateParams |
-  ResetChatParams |
-  ForwardTokensAndSampleParams |
-  ChatCompletionNonStreamingParams |
-  ChatCompletionStreamInitParams |
-  CustomRequestParams |
-  InitProgressReport |
-  string |
-  null |
-  number |
-  ChatCompletion |
-  ChatCompletionChunk |
-  void;
+  | GenerateProgressCallbackParams
+  | ReloadParams
+  | GenerateParams
+  | ResetChatParams
+  | ForwardTokensAndSampleParams
+  | ChatCompletionNonStreamingParams
+  | ChatCompletionStreamInitParams
+  | CustomRequestParams
+  | InitProgressReport
+  | string
+  | null
+  | number
+  | ChatCompletion
+  | ChatCompletionChunk
+  | void;
 /**
  * The message used in exchange between worker
  * and the main thread.
  */
 
-export interface WorkerMessage {
+export type WorkerRequest = {
   kind: RequestKind;
   uuid: string;
   content: MessageContent;
-}
+};
+
+export type OneTimeWorkerResponse = {
+  kind: "return" | "throw";
+  uuid: string;
+  content: MessageContent;
+};
+
+export type InitProgressWorkerResponse = {
+  kind: "initProgressCallback";
+  uuid: string;
+  content: InitProgressReport;
+};
+
+export type GenerateProgressWorkerResponse = {
+  kind: "generateProgressCallback";
+  uuid: string;
+  content: {
+    step: number;
+    currentMessage: string;
+  };
+};
+
+export type WorkerResponse =
+  | OneTimeWorkerResponse
+  | InitProgressWorkerResponse
+  | GenerateProgressWorkerResponse;

src/message.ts

@@ -1,45 +1,19 @@
 import * as tvmjs from "tvmjs";
-import { AppConfig, ChatOptions, EngineConfig } from "./config";
-import { ReloadParams, WorkerMessage } from "./message";
-import { EngineInterface } from "./types";
-import {
-  ChatWorker,
-  EngineWorkerHandler,
-  WebWorkerEngine,
-  PostMessageHandler,
-} from "./web_worker";
+import { AppConfig, ChatOptions, EngineConfig, ModelRecord } from "./config";
+import { ReloadParams, WorkerRequest, WorkerResponse } from "./message";
+import { EngineInterface, InitProgressReport } from "./types";
+import { EngineWorkerHandler, WebWorkerEngine, ChatWorker } from "./web_worker";
 import { areAppConfigsEqual, areChatOptionsEqual } from "./utils";
 
-/**
- * A post message handler that sends messages to a chrome.runtime.Port.
- */
-export class PortPostMessageHandler implements PostMessageHandler {
-  port: chrome.runtime.Port;
-  enabled: boolean = true;
-
-  constructor(port: chrome.runtime.Port) {
-    this.port = port;
-  }
-
-  /**
-   * Close the PortPostMessageHandler. This will prevent any further messages
-   */
-  close() {
-    this.enabled = false;
-  }
+/* Service Worker Script */
 
-  postMessage(event: any) {
-    if (this.enabled) {
-      this.port.postMessage(event);
-    }
-  }
-}
+type IServiceWorker = globalThis.ServiceWorker;
 
 /**
  * Worker handler that can be used in a ServiceWorker.
- * 
+ *
  * @example
- * 
+ *
  * const engine = new Engine();
  * let handler;
  * chrome.runtime.onConnect.addListener(function (port) {
@@ -56,29 +30,74 @@ export class ServiceWorkerEngineHandler extends EngineWorkerHandler {
   chatOpts?: ChatOptions;
   appConfig?: AppConfig;
 
-  constructor(engine: EngineInterface, port: chrome.runtime.Port) {
-    let portHandler = new PortPostMessageHandler(port);
-    super(engine, portHandler);
+  private clientRegistry = new Map<
+    string,
+    IServiceWorker | Client | MessagePort
+  >();
+  private initReuqestUuid?: string;
 
-    port.onDisconnect.addListener(() => {
-      portHandler.close();
-    });
-  }
+  constructor(engine: EngineInterface) {
+    if (!self || !("addEventListener" in self)) {
+      throw new Error(
+        "ServiceWorkerGlobalScope is not defined. ServiceWorkerEngineHandler must be created in service worker script."
+      );
+    }
+    const postMessageHandler = {
+      postMessage: (message: WorkerResponse) => {
+        if (this.clientRegistry.has(message.uuid)) {
+          const client = this.clientRegistry.get(message.uuid);
+          client?.postMessage(message);
 
-  setPort(port: chrome.runtime.Port) {
-    let portHandler = new PortPostMessageHandler(port);
-    this.setPostMessageHandler(portHandler);
-    port.onDisconnect.addListener(() => {
-      portHandler.close();
+          if (message.kind === "return" || message.kind === "throw") {
+            this.clientRegistry.delete(message.uuid);
+          } else {
+            // TODO: Delete clientRegistry after complete to avoid memory leak?
+          }
+        }
+      },
+    };
+    const initProgressCallback = (report: InitProgressReport) => {
+      const msg: WorkerResponse = {
+        kind: "initProgressCallback",
+        uuid: this.initReuqestUuid || "",
+        content: report,
+      };
+      this.postMessageInternal(msg);
+    };
+    super(engine, postMessageHandler, initProgressCallback);
+    const onmessage = this.onmessage.bind(this);
+
+    self.addEventListener("message", (event) => {
+      const message = event as unknown as ExtendableMessageEvent;
+      if (message.source) {
+        this.clientRegistry.set(message.data.uuid, message.source);
+      }
+      message.waitUntil(
+        new Promise((resolve, reject) => {
+          onmessage(message, resolve, reject);
+        })
+      );
     });
   }
 
-  onmessage(event: any): void {
-    if (event.type === "keepAlive") {
+  onmessage(
+    event: ExtendableMessageEvent,
+    onComplete?: (value: any) => void,
+    onError?: () => void
+  ): void {
+    const msg = event.data as WorkerRequest;
+
+    if (msg.kind === "keepAlive") {
+      const reply: WorkerRequest = {
+        kind: "heartbeat",
+        uuid: msg.uuid,
+        content: "",
+      };
+      this.postMessageInternal(reply);
+      onComplete?.(reply);
       return;
     }
 
-    const msg = event as WorkerMessage;
     if (msg.kind === "init") {
       this.handleTask(msg.uuid, async () => {
         const params = msg.content as ReloadParams;
@@ -104,9 +123,11 @@ export class ServiceWorkerEngineHandler extends EngineWorkerHandler {
             timeElapsed: 0,
             text: "Finish loading on " + gpuLabel,
           });
+          onComplete?.(null);
           return null;
         }
 
+        this.initReuqestUuid = msg.uuid;
         await this.engine.reload(
           params.modelId,
           params.chatOpts,
@@ -115,31 +136,60 @@ export class ServiceWorkerEngineHandler extends EngineWorkerHandler {
         this.modelId = params.modelId;
         this.chatOpts = params.chatOpts;
         this.appConfig = params.appConfig;
+        onComplete?.(null);
         return null;
       });
       return;
     }
-    super.onmessage(event);
+    super.onmessage(msg, onComplete, onError);
+  }
+}
+
+/* Webapp Client */
+/**
+ * PostMessageHandler wrapper for sending message from client to service worker
+ */
+export class ServiceWorker implements ChatWorker {
+  serviceWorker: IServiceWorker;
+
+  constructor(serviceWorker: IServiceWorker) {
+    this.serviceWorker = serviceWorker;
+  }
+
+  // ServiceWorkerEngine will later overwrite this
+  onmessage() {}
+
+  postMessage(message: WorkerRequest) {
+    if (!("serviceWorker" in navigator)) {
+      throw new Error("Service worker API is not available");
+    }
+    const serviceWorker = (navigator.serviceWorker as ServiceWorkerContainer)
+      .controller;
+    if (!serviceWorker) {
+      throw new Error("There is no active service worker");
+    }
+    serviceWorker.postMessage(message);
   }
 }
 
 /**
  * Create a ServiceWorkerEngine.
- * 
+ *
  * @param modelId The model to load, needs to either be in `webllm.prebuiltAppConfig`, or in
  * `engineConfig.appConfig`.
  * @param engineConfig Optionally configures the engine, see `webllm.EngineConfig` for more.
- * @param keepAliveMs The interval to send keep alive messages to the service worker.
- * See [Service worker lifecycle](https://developer.chrome.com/docs/extensions/develop/concepts/service-workers/lifecycle#idle-shutdown)
- * The default is 10s.
  * @returns An initialized `WebLLM.ServiceWorkerEngine` with `modelId` loaded.
  */
 export async function CreateServiceWorkerEngine(
   modelId: string,
-  engineConfig?: EngineConfig,
-  keepAliveMs: number = 10000
+  engineConfig?: EngineConfig
 ): Promise<ServiceWorkerEngine> {
-  const serviceWorkerEngine = new ServiceWorkerEngine(keepAliveMs);
+  if (!("serviceWorker" in navigator)) {
+    throw new Error("Service worker API is not available");
+  }
+  const registration = await (navigator.serviceWorker as ServiceWorkerContainer)
+    .ready;
+  const serviceWorkerEngine = new ServiceWorkerEngine(registration.active!);
   serviceWorkerEngine.setInitProgressCallback(
     engineConfig?.initProgressCallback
   );
@@ -151,57 +201,44 @@ export async function CreateServiceWorkerEngine(
   return serviceWorkerEngine;
 }
 
-class PortAdapter implements ChatWorker {
-  port: chrome.runtime.Port;
-  private _onmessage!: (message: any) => void;
-
-  constructor(port: chrome.runtime.Port) {
-    this.port = port;
-    this.port.onMessage.addListener(this.handleMessage.bind(this));
-  }
-
-  // Wrapper to handle incoming messages and delegate to onmessage if available
-  private handleMessage(message: any) {
-    if (this._onmessage) {
-      this._onmessage(message);
-    }
-  }
-
-  // Getter and setter for onmessage to manage adding/removing listeners
-  get onmessage(): (message: any) => void {
-    return this._onmessage;
-  }
-
-  set onmessage(listener: (message: any) => void) {
-    this._onmessage = listener;
-  }
-
-  // Wrap port.postMessage to maintain 'this' context
-  postMessage = (message: any): void => {
-    this.port.postMessage(message);
-  };
-}
-
 /**
  * A client of Engine that exposes the same interface
  */
 export class ServiceWorkerEngine extends WebWorkerEngine {
-  port: chrome.runtime.Port;
+  missedHeatbeat = 0;
+
+  constructor(worker: IServiceWorker, keepAliveMs = 10000) {
+    if (!("serviceWorker" in navigator)) {
+      throw new Error("Service worker API is not available");
+    }
+    super(new ServiceWorker(worker));
+    const onmessage = this.onmessage.bind(this);
+
+    (navigator.serviceWorker as ServiceWorkerContainer).addEventListener(
+      "message",
+      (event: MessageEvent) => {
+        const msg = event.data;
+        try {
+          if (msg.kind === "heartbeat") {
+            this.missedHeatbeat = 0;
+            return;
+          }
+          onmessage(msg);
+        } 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);
+          }
+        }
+      }
+    );
 
-  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();
+      this.worker.postMessage({ kind: "keepAlive", uuid: crypto.randomUUID() });
+      this.missedHeatbeat += 1;
     }, keepAliveMs);
   }
 
-  keepAlive() {
-    this.worker.postMessage({ kind: "keepAlive" });
-  }
-
   /**
    * Initialize the chat with a model.
    *
@@ -218,7 +255,7 @@ export class ServiceWorkerEngine extends WebWorkerEngine {
     chatOpts?: ChatOptions,
     appConfig?: AppConfig
   ): Promise<void> {
-    const msg: WorkerMessage = {
+    const msg: WorkerRequest = {
       kind: "init",
       uuid: crypto.randomUUID(),
       content: {

src/service_worker.ts

@@ -7,9 +7,10 @@
     "sourceMap": true,
     "strict": true,
     "moduleResolution": "Node",
-    "esModuleInterop": true
+    "esModuleInterop": true,
+    "lib": ["dom", "WebWorker"]
   },
-  "typeRoots": [ "./node_modules/@webgpu/types", "./node_modules/@types"],
+  "typeRoots": ["./node_modules/@webgpu/types", "./node_modules/@types"],
   "include": ["src"],
   "exclude": ["node_modules", "build", "dist", "rollup.config.cjs"]
 }