use-browser-llm 0.1.0 → 0.1.2

package/dist/index.js

@@ -1,5 +1,12 @@
 // src/use-browser-llm.ts
-import { useCallback, useEffect, useReducer, useRef, useState } from "react";
+import {
+  useCallback,
+  useEffect,
+  useLayoutEffect,
+  useReducer,
+  useRef,
+  useState
+} from "react";
 
 // src/engine-client.ts
 import * as Comlink from "comlink";
@@ -98,41 +105,58 @@ async function detectWebGPUSupport() {
 }
 
 // src/to-async-generator.ts
-async function* toAsyncGenerator(run) {
+async function* toAsyncGenerator(run, signal) {
   const queue = [];
   let resolveNext = null;
   let done = false;
   let error = null;
+  const onAbort = () => {
+    if (!done) {
+      error = signal?.reason ?? new Error("aborted");
+      done = true;
+      resolveNext?.();
+      resolveNext = null;
+    }
+  };
+  signal?.addEventListener("abort", onAbort);
   run((token) => {
     queue.push(token);
     resolveNext?.();
     resolveNext = null;
   }).then(
     () => {
-      done = true;
-      resolveNext?.();
-      resolveNext = null;
+      if (!done) {
+        done = true;
+        resolveNext?.();
+        resolveNext = null;
+      }
     },
     (err) => {
-      error = err;
-      done = true;
-      resolveNext?.();
-      resolveNext = null;
+      if (!done) {
+        error = err;
+        done = true;
+        resolveNext?.();
+        resolveNext = null;
+      }
     }
   );
-  while (true) {
-    while (queue.length > 0) {
-      yield queue.shift();
-    }
-    if (done) {
-      if (error) {
-        throw error;
+  try {
+    while (true) {
+      while (queue.length > 0) {
+        yield queue.shift();
       }
-      return;
+      if (done) {
+        if (error) {
+          throw error;
+        }
+        return;
+      }
+      await new Promise((resolve) => {
+        resolveNext = resolve;
+      });
     }
-    await new Promise((resolve) => {
-      resolveNext = resolve;
-    });
+  } finally {
+    signal?.removeEventListener("abort", onAbort);
   }
 }
 
@@ -144,6 +168,8 @@ var initialState = {
   cacheStatus: "idle"
 };
 var LOAD_INACTIVITY_TIMEOUT_MS = 3e4;
+var GENERATE_TIMEOUT_MS = 9e4;
+var STREAM_INACTIVITY_TIMEOUT_MS = 3e4;
 function reducer(state, action) {
   switch (action.type) {
     case "reset":
@@ -179,7 +205,9 @@ function useBrowserLLM(modelId) {
   const [state, dispatch] = useReducer(reducer, initialState);
   const clientRef = useRef(null);
   const statusRef = useRef(state.status);
-  statusRef.current = state.status;
+  useLayoutEffect(() => {
+    statusRef.current = state.status;
+  }, [state.status]);
   const isGeneratingRef = useRef(false);
   const [isGenerating, setIsGenerating] = useState(false);
   const [generationError, setGenerationError] = useState(null);
@@ -233,6 +261,8 @@ function useBrowserLLM(modelId) {
       unsubscribeCrash = client.onCrash((error) => {
         if (!cancelled) {
           clearWatchdog();
+          clientRef.current?.terminate();
+          clientRef.current = null;
           dispatch({ type: "error", error });
         }
       });
@@ -281,16 +311,63 @@ function useBrowserLLM(modelId) {
       if (isGeneratingRef.current) {
         throw new HookBusyError();
       }
+      const client = clientRef.current;
       isGeneratingRef.current = true;
       setIsGenerating(true);
       setGenerationError(null);
+      let settled = false;
+      let timeoutId = null;
+      const crashSub = {
+        unsubscribe: null
+      };
       try {
-        return await clientRef.current.generate(messages);
+        return await new Promise((resolve, reject) => {
+          crashSub.unsubscribe = client.onCrash((error) => {
+            if (settled) {
+              return;
+            }
+            settled = true;
+            reject(error);
+          });
+          timeoutId = setTimeout(() => {
+            if (settled) {
+              return;
+            }
+            settled = true;
+            const error = new WorkerCrashError(
+              `no response for ${GENERATE_TIMEOUT_MS / 1e3}s \u2014 the worker may have been killed by the browser (e.g. out of memory)`
+            );
+            if (clientRef.current === client) {
+              clientRef.current = null;
+              client.terminate();
+              dispatch({ type: "error", error });
+            }
+            reject(error);
+          }, GENERATE_TIMEOUT_MS);
+          client.generate(messages).then(
+            (result) => {
+              if (!settled) {
+                settled = true;
+                resolve(result);
+              }
+            },
+            (err) => {
+              if (!settled) {
+                settled = true;
+                reject(err instanceof Error ? err : new Error(String(err)));
+              }
+            }
+          );
+        });
       } catch (err) {
         const error = err instanceof Error ? err : new Error(String(err));
         setGenerationError(error);
         throw error;
       } finally {
+        if (timeoutId !== null) {
+          clearTimeout(timeoutId);
+        }
+        crashSub.unsubscribe?.();
         isGeneratingRef.current = false;
         setIsGenerating(false);
       }
@@ -304,6 +381,7 @@ function useBrowserLLM(modelId) {
       isGeneratingRef,
       setIsGenerating,
       setGenerationError,
+      dispatch,
       messages
     ),
     []
@@ -321,7 +399,7 @@ function useBrowserLLM(modelId) {
     abort
   };
 }
-async function* streamGenerateImpl(statusRef, clientRef, isGeneratingRef, setIsGenerating, setGenerationError, messages) {
+async function* streamGenerateImpl(statusRef, clientRef, isGeneratingRef, setIsGenerating, setGenerationError, dispatch, messages) {
   if (statusRef.current !== "ready" || !clientRef.current) {
     throw new HookNotReadyError();
   }
@@ -332,13 +410,48 @@ async function* streamGenerateImpl(statusRef, clientRef, isGeneratingRef, setIsG
   isGeneratingRef.current = true;
   setIsGenerating(true);
   setGenerationError(null);
+  const controller = new AbortController();
+  let watchdog = null;
+  const clearWatchdog = () => {
+    if (watchdog !== null) {
+      clearTimeout(watchdog);
+      watchdog = null;
+    }
+  };
+  const resetWatchdog = () => {
+    clearWatchdog();
+    watchdog = setTimeout(() => {
+      const error = new WorkerCrashError(
+        `no token for ${STREAM_INACTIVITY_TIMEOUT_MS / 1e3}s \u2014 the worker may have been killed by the browser (e.g. out of memory)`
+      );
+      controller.abort(error);
+      if (clientRef.current === client) {
+        clientRef.current = null;
+        client.terminate();
+        dispatch({ type: "error", error });
+      }
+    }, STREAM_INACTIVITY_TIMEOUT_MS);
+  };
+  const unsubscribeCrash = client.onCrash((error) => {
+    controller.abort(error);
+  });
   try {
-    yield* toAsyncGenerator((onToken) => client.streamGenerate(messages, onToken));
+    resetWatchdog();
+    const tokens = toAsyncGenerator(
+      (onToken) => client.streamGenerate(messages, onToken),
+      controller.signal
+    );
+    for await (const token of tokens) {
+      resetWatchdog();
+      yield token;
+    }
   } catch (err) {
     const error = err instanceof Error ? err : new Error(String(err));
     setGenerationError(error);
     throw error;
   } finally {
+    clearWatchdog();
+    unsubscribeCrash();
     isGeneratingRef.current = false;
     setIsGenerating(false);
     client.abort().catch(() => {

package/package.json

@@ -1,7 +1,19 @@
 {
   "name": "use-browser-llm",
-  "version": "0.1.0",
+  "version": "0.1.2",
   "description": "Headless React hook for running LLMs locally in the browser via WebGPU, with zero backend and complete privacy.",
+  "keywords": [
+    "react",
+    "react-hooks",
+    "llm",
+    "ai",
+    "webgpu",
+    "web-llm",
+    "browser",
+    "on-device-ai",
+    "privacy",
+    "offline"
+  ],
   "license": "MIT",
   "repository": {
     "type": "git",