feat: implemented the v1 browser embedded LLM support (Gemma 4 E4B based)

2026-05-14 18:55:23 -07:00
parent 0e6e2736d0
commit 98c2b97413
30 changed files with 19462 additions and 153 deletions
@@ -0,0 +1,42 @@
+export const LOCAL_LLM_PROVIDER_KEY = 'local_browser';
+export const LOCAL_LLM_MODEL_URL =
+  'http://localhost:3000/models/gemma-4-E4B-it-web.task';
+export const LOCAL_LLM_MODEL_FILENAME = 'gemma-4-E4B-it-web.task';
+export const LOCAL_LLM_DISPLAY_NAME = 'Gemma 4 E4B';
+export const LOCAL_LLM_LEGACY_FILENAMES = [
+  'gemma-3n-E4B-it-int4-Web.litertlm',
+];
+
+export interface LocalLLMRuntimeSupport {
+  supported: boolean;
+  webgpuExposed: boolean;
+  crossOriginIsolated: boolean;
+  sharedArrayBufferAvailable: boolean;
+  secureContext: boolean;
+  reason: string | null;
+}
+
+export function detectLocalLLMRuntimeSupport(): LocalLLMRuntimeSupport {
+  const secureContext = typeof window !== 'undefined' ? window.isSecureContext : false;
+  const crossOriginIsolated = typeof window !== 'undefined' ? window.crossOriginIsolated : false;
+  const sharedArrayBufferAvailable = typeof SharedArrayBuffer !== 'undefined';
+  const webgpuExposed = typeof navigator !== 'undefined' && 'gpu' in navigator;
+
+  let reason: string | null = null;
+  if (!secureContext) {
+    reason = 'Local browser LLM requires a secure context (HTTPS or localhost).';
+  } else if (!crossOriginIsolated || !sharedArrayBufferAvailable) {
+    reason = 'Local browser LLM requires SharedArrayBuffer support. Ensure COOP/COEP headers are enabled.';
+  } else if (!webgpuExposed) {
+    reason = 'Local browser LLM currently requires a browser with WebGPU support.';
+  }
+
+  return {
+    supported: reason === null,
+    webgpuExposed,
+    crossOriginIsolated,
+    sharedArrayBufferAvailable,
+    secureContext,
+    reason,
+  };
+}
@@ -0,0 +1,92 @@
+import { OpfsModelCache, type ModelDownloadProgress } from './opfsModelCache';
+import { LOCAL_LLM_MODEL_FILENAME } from './localLLMConfig';
+
+const cache = new OpfsModelCache({ directoryName: 'models' });
+let writingToCachePromise: Promise<void> | null = null;
+
+export { type ModelDownloadProgress };
+
+export interface CachedModelStreamResult {
+  reader: ReadableStreamDefaultReader<Uint8Array>;
+  totalBytes: number;
+  fromCache: boolean;
+  cacheWritePromise: Promise<void> | null;
+}
+
+export class LocalLLMModelCache {
+  public static async exists(filename: string = LOCAL_LLM_MODEL_FILENAME): Promise<boolean> {
+    return cache.exists(filename);
+  }
+
+  public static async getFile(filename: string = LOCAL_LLM_MODEL_FILENAME): Promise<File> {
+    return cache.getFile(filename);
+  }
+
+  public static async getArrayBuffer(filename: string = LOCAL_LLM_MODEL_FILENAME): Promise<ArrayBuffer> {
+    return cache.getArrayBuffer(filename);
+  }
+
+  public static async delete(filename: string = LOCAL_LLM_MODEL_FILENAME): Promise<void> {
+    await cache.delete(filename);
+  }
+
+  public static async loadModelReaderWithCache(
+    sourceUrl: string,
+    filename: string = LOCAL_LLM_MODEL_FILENAME,
+    onProgress?: (progress: ModelDownloadProgress & { fromCache: boolean }) => void,
+  ): Promise<CachedModelStreamResult> {
+    if (writingToCachePromise) {
+      await writingToCachePromise.catch(() => {});
+    }
+
+    if (await this.exists(filename)) {
+      const file = await this.getFile(filename);
+      onProgress?.({
+        receivedBytes: file.size,
+        totalBytes: file.size,
+        percent: 100,
+        fromCache: true,
+      });
+      return {
+        reader: file.stream().getReader(),
+        totalBytes: file.size,
+        fromCache: true,
+        cacheWritePromise: null,
+      };
+    }
+
+    const response = await fetch(sourceUrl);
+    if (!response.ok || !response.body) {
+      throw new Error(`Model download failed (${response.status})`);
+    }
+
+    const totalBytesHeader = response.headers.get('Content-Length');
+    const totalBytes = totalBytesHeader ? Number(totalBytesHeader) : 0;
+    const [streamForConsumer, streamForCache] = response.body.tee();
+
+    writingToCachePromise = cache.downloadStream(
+      streamForCache,
+      filename,
+      totalBytes > 0 ? totalBytes : null,
+      progress => onProgress?.({ ...progress, fromCache: false }),
+    );
+    writingToCachePromise = writingToCachePromise.finally(() => {
+      writingToCachePromise = null;
+    });
+
+    return {
+      reader: streamForConsumer.getReader(),
+      totalBytes,
+      fromCache: false,
+      cacheWritePromise: writingToCachePromise,
+    };
+  }
+
+  public static async download(
+    sourceUrl: string,
+    filename: string = LOCAL_LLM_MODEL_FILENAME,
+    onProgress?: (progress: ModelDownloadProgress) => void,
+  ): Promise<void> {
+    await cache.download(sourceUrl, filename, onProgress);
+  }
+}
@@ -0,0 +1,159 @@
+import {
+  detectLocalLLMRuntimeSupport,
+  LOCAL_LLM_LEGACY_FILENAMES,
+  LOCAL_LLM_MODEL_FILENAME,
+  LOCAL_LLM_MODEL_URL,
+  type LocalLLMRuntimeSupport,
+} from './localLLMConfig';
+import { LocalLLMModelCache } from './localLLMModelCache';
+
+export interface LocalLLMModelState {
+  isCached: boolean;
+  isChecking: boolean;
+  isDownloading: boolean;
+  isDeleting: boolean;
+  progressPercent: number;
+  progressText: string;
+  error: string;
+  runtimeSupport: LocalLLMRuntimeSupport;
+}
+
+type Listener = (state: LocalLLMModelState) => void;
+
+export class LocalLLMModelManager {
+  private static listeners = new Set<Listener>();
+  private static initialized = false;
+  private static state: LocalLLMModelState = {
+    isCached: false,
+    isChecking: false,
+    isDownloading: false,
+    isDeleting: false,
+    progressPercent: 0,
+    progressText: '',
+    error: '',
+    runtimeSupport: detectLocalLLMRuntimeSupport(),
+  };
+
+  public static subscribe(listener: Listener): () => void {
+    this.listeners.add(listener);
+    listener(this.getState());
+    if (!this.initialized) {
+      this.initialized = true;
+      void this.refresh();
+    }
+    return () => this.listeners.delete(listener);
+  }
+
+  public static getState(): LocalLLMModelState {
+    return { ...this.state, runtimeSupport: { ...this.state.runtimeSupport } };
+  }
+
+  public static async refresh(): Promise<void> {
+    this.setState({
+      isChecking: true,
+      runtimeSupport: detectLocalLLMRuntimeSupport(),
+    });
+    try {
+      await this.cleanupLegacyEntries();
+      const isCached = await LocalLLMModelCache.exists();
+      this.setState({ isCached, error: '' });
+    } catch (error) {
+      this.setState({ error: error instanceof Error ? error.message : String(error) });
+    } finally {
+      this.setState({ isChecking: false });
+    }
+  }
+
+  public static async ensureRuntimeSupported(): Promise<void> {
+    const runtimeSupport = detectLocalLLMRuntimeSupport();
+    this.setState({ runtimeSupport });
+    if (!runtimeSupport.supported) {
+      throw new Error(runtimeSupport.reason ?? 'Local browser LLM is not supported in this browser.');
+    }
+
+    await this.cleanupLegacyEntries();
+  }
+
+  public static async deleteCachedModel(): Promise<void> {
+    this.setState({ isDeleting: true, error: '' });
+    try {
+      await LocalLLMModelCache.delete();
+      await this.cleanupLegacyEntries();
+      this.setState({
+        isCached: false,
+        progressPercent: 0,
+        progressText: '',
+      });
+    } catch (error) {
+      this.setState({ error: error instanceof Error ? error.message : String(error) });
+      throw error;
+    } finally {
+      this.setState({ isDeleting: false });
+    }
+  }
+
+  private static setState(partial: Partial<LocalLLMModelState>): void {
+    this.state = {
+      ...this.state,
+      ...partial,
+    };
+    for (const listener of this.listeners) {
+      listener(this.getState());
+    }
+  }
+
+  private static async cleanupLegacyEntries(): Promise<void> {
+    await Promise.all(
+      LOCAL_LLM_LEGACY_FILENAMES.map(async legacyFilename => {
+        try {
+          await LocalLLMModelCache.delete(legacyFilename);
+        } catch {
+          // Ignore best-effort legacy cleanup failures.
+        }
+      }),
+    );
+  }
+
+  public static notifyLoadStart(fromCache: boolean): void {
+    this.setState({
+      isDownloading: !fromCache,
+      progressPercent: fromCache ? 100 : 0,
+      progressText: fromCache ? 'Loading local language model from browser cache...' : 'Downloading local language model...',
+      error: '',
+    });
+  }
+
+  public static notifyLoadProgress(receivedBytes: number, totalBytes: number | null, fromCache: boolean): void {
+    const receivedMb = (receivedBytes / (1024 * 1024)).toFixed(1);
+    const totalMb = totalBytes ? (totalBytes / (1024 * 1024)).toFixed(1) : null;
+    this.setState({
+      isDownloading: !fromCache,
+      progressPercent: totalBytes ? (receivedBytes / totalBytes) * 100 : 0,
+      progressText: fromCache
+        ? 'Loading local language model from browser cache...'
+        : totalMb
+          ? `Downloading local language model... ${receivedMb} / ${totalMb} MB`
+          : `Downloading local language model... ${receivedMb} MB`,
+      error: '',
+    });
+  }
+
+  public static notifyCacheReady(): void {
+    this.setState({
+      isCached: true,
+      isDownloading: false,
+      progressPercent: 100,
+      progressText: 'Local language model is ready.',
+      error: '',
+    });
+  }
+
+  public static notifyLoadError(error: unknown): void {
+    this.setState({
+      isDownloading: false,
+      progressPercent: 0,
+      progressText: '',
+      error: error instanceof Error ? error.message : String(error),
+    });
+  }
+}
@@ -1,50 +1,25 @@
 import { LOCAL_SEPARATOR_MODEL_FILENAME } from './localSeparatorConfig';
+import { OpfsModelCache, type ModelDownloadProgress } from './opfsModelCache';

-export interface ModelDownloadProgress {
-  receivedBytes: number;
-  totalBytes: number | null;
-  percent: number;
-}
+const cache = new OpfsModelCache({ directoryName: 'models' });
+
+export { type ModelDownloadProgress };

 export class LocalSeparatorModelCache {
-  private static readonly MODELS_DIR = 'models';
-  private static readonly TEMP_SUFFIX = '.download';
-
  public static async exists(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise<boolean> {
-    try {
-      const modelsDir = await this.getModelsDir();
-      await modelsDir.getFileHandle(filename);
-      return true;
-    } catch {
-      return false;
-    }
+    return cache.exists(filename);
  }

  public static async getFile(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise<File> {
-    const modelsDir = await this.getModelsDir();
-    const fileHandle = await modelsDir.getFileHandle(filename);
-    return fileHandle.getFile();
+    return cache.getFile(filename);
  }

  public static async getArrayBuffer(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise<ArrayBuffer> {
-    const file = await this.getFile(filename);
-    return file.arrayBuffer();
+    return cache.getArrayBuffer(filename);
  }

  public static async delete(filename: string = LOCAL_SEPARATOR_MODEL_FILENAME): Promise<void> {
-    try {
-      const modelsDir = await this.getModelsDir();
-      await modelsDir.removeEntry(filename);
-    } catch {
-      // Ignore missing file cleanup.
-    }
-
-    try {
-      const modelsDir = await this.getModelsDir();
-      await modelsDir.removeEntry(`${filename}${this.TEMP_SUFFIX}`);
-    } catch {
-      // Ignore missing temp file cleanup.
-    }
+    await cache.delete(filename);
  }

  public static async download(
@@ -52,75 +27,6 @@ export class LocalSeparatorModelCache {
    filename: string = LOCAL_SEPARATOR_MODEL_FILENAME,
    onProgress?: (progress: ModelDownloadProgress) => void,
  ): Promise<void> {
-    const response = await fetch(sourceUrl);
-    if (!response.ok) {
-      throw new Error(`Model download failed (${response.status})`);
-    }
-
-    const modelsDir = await this.getModelsDir();
-    const tempName = `${filename}${this.TEMP_SUFFIX}`;
-    await this.delete(filename);
-
-    const tempHandle = await modelsDir.getFileHandle(tempName, { create: true });
-    const writable = await tempHandle.createWritable();
-
-    try {
-      const totalBytesHeader = response.headers.get('Content-Length');
-      const totalBytes = totalBytesHeader ? Number(totalBytesHeader) : null;
-
-      if (!response.body) {
-        const buffer = await response.arrayBuffer();
-        await writable.write(buffer);
-        onProgress?.({
-          receivedBytes: buffer.byteLength,
-          totalBytes,
-          percent: 100,
-        });
-      } else {
-        const reader = response.body.getReader();
-        let receivedBytes = 0;
-
-        while (true) {
-          const { done, value } = await reader.read();
-          if (done) break;
-          if (!value) continue;
-          await writable.write(value);
-          receivedBytes += value.byteLength;
-          onProgress?.({
-            receivedBytes,
-            totalBytes,
-            percent: totalBytes ? (receivedBytes / totalBytes) * 100 : 0,
-          });
-        }
-      }
-    } catch (error) {
-      await writable.abort();
-      await this.delete(filename);
-      throw error;
-    }
-
-    await writable.close();
-
-    const finalHandle = await modelsDir.getFileHandle(filename, { create: true });
-    const finalWritable = await finalHandle.createWritable();
-    try {
-      const tempFile = await tempHandle.getFile();
-      await finalWritable.write(await tempFile.arrayBuffer());
-      await finalWritable.close();
-    } catch (error) {
-      await finalWritable.abort();
-      throw error;
-    } finally {
-      try {
-        await modelsDir.removeEntry(tempName);
-      } catch {
-        // Ignore temp cleanup errors.
-      }
-    }
-  }
-
-  private static async getModelsDir(): Promise<FileSystemDirectoryHandle> {
-    const root = await navigator.storage.getDirectory();
-    return root.getDirectoryHandle(this.MODELS_DIR, { create: true });
+    await cache.download(sourceUrl, filename, onProgress);
  }
 }
@@ -2,6 +2,7 @@ import { clearChatHistoryAndUI } from '../chatUtil';
 import { useProjectStore } from '../../stores/projectStore';
 import { ConfigManager } from '../../core/config/ConfigManager';
 import { SystemPrompts } from '../../agent/core/SystemPrompts';
+import { detectLocalLLMRuntimeSupport, LOCAL_LLM_PROVIDER_KEY } from '../localLLMConfig';

 export interface UserMessageFilterResult {
  // Whether to render the user message bubble (div.message-user)
@@ -134,7 +135,18 @@ export async function processUserMessage(originalMessage: string): Promise<UserM
      }
      const provider = (configManager.get('general.llm_provider') as string) || 'openai';

-      if (provider === 'openai') {
+      if (provider === LOCAL_LLM_PROVIDER_KEY) {
+        const runtimeSupport = detectLocalLLMRuntimeSupport();
+        if (!runtimeSupport.supported) {
+          return {
+            displayUserMessage: true,
+            sendToLLM: false,
+            finalMessageForLLM: null,
+            pseudoAssistantResponse: runtimeSupport.reason ?? 'Local browser LLM is not supported in this environment.',
+            metadata: { error: 'local_browser_unsupported' }
+          };
+        }
+      } else if (provider === 'openai') {
        const openaiKey = (configManager.get('general.openai.api_key') as string) || '';
        if (openaiKey.trim() === '') {
          const url = `${import.meta.env.BASE_URL}chat/error_no_openai_key.md`;
@@ -256,4 +268,3 @@ export async function processUserMessage(originalMessage: string): Promise<UserM
  }
 }

-
@@ -0,0 +1,188 @@
+export interface ModelDownloadProgress {
+  receivedBytes: number;
+  totalBytes: number | null;
+  percent: number;
+}
+
+interface OpfsModelCacheOptions {
+  directoryName?: string;
+  sizeSuffix?: string;
+  tempSuffix?: string;
+}
+
+export class OpfsModelCache {
+  private readonly directoryName: string;
+  private readonly sizeSuffix: string;
+  private readonly tempSuffix: string;
+
+  constructor(options: OpfsModelCacheOptions = {}) {
+    this.directoryName = options.directoryName ?? 'models';
+    this.sizeSuffix = options.sizeSuffix ?? '.size';
+    this.tempSuffix = options.tempSuffix ?? '.download';
+  }
+
+  public async exists(filename: string): Promise<boolean> {
+    try {
+      const dir = await this.getDir();
+      const fileHandle = await dir.getFileHandle(filename);
+      const sizeHandle = await dir.getFileHandle(this.getSizeFilename(filename));
+      const [file, sizeFile] = await Promise.all([fileHandle.getFile(), sizeHandle.getFile()]);
+      const expectedSize = Number(await sizeFile.text());
+      if (!Number.isFinite(expectedSize) || expectedSize <= 0) {
+        await this.delete(filename);
+        return false;
+      }
+      if (file.size !== expectedSize) {
+        await this.delete(filename);
+        return false;
+      }
+      return true;
+    } catch {
+      return false;
+    }
+  }
+
+  public async getFile(filename: string): Promise<File> {
+    const dir = await this.getDir();
+    const handle = await dir.getFileHandle(filename);
+    const file = await handle.getFile();
+    console.log('[opfsModelCache] Opened cached file.', {
+      filename,
+      size: file.size,
+    });
+    return file;
+  }
+
+  public async getArrayBuffer(filename: string): Promise<ArrayBuffer> {
+    const file = await this.getFile(filename);
+    return file.arrayBuffer();
+  }
+
+  public async delete(filename: string): Promise<void> {
+    const dir = await this.getDir();
+    await this.removeIfExists(dir, filename);
+    await this.removeIfExists(dir, this.getSizeFilename(filename));
+    await this.removeIfExists(dir, `${filename}${this.tempSuffix}`);
+    await this.removeIfExists(dir, `${this.getSizeFilename(filename)}${this.tempSuffix}`);
+  }
+
+  public async download(
+    sourceUrl: string,
+    filename: string,
+    onProgress?: (progress: ModelDownloadProgress) => void,
+  ): Promise<void> {
+    const response = await fetch(sourceUrl);
+    if (!response.ok) {
+      throw new Error(`Model download failed (${response.status})`);
+    }
+    const totalBytesHeader = response.headers.get('Content-Length');
+    const totalBytes = totalBytesHeader ? Number(totalBytesHeader) : null;
+    if (!response.body) {
+      throw new Error('Model download response did not include a readable body.');
+    }
+    await this.downloadStream(response.body, filename, totalBytes, onProgress);
+  }
+
+  public async downloadStream(
+    stream: ReadableStream<Uint8Array>,
+    filename: string,
+    totalBytes: number | null,
+    onProgress?: (progress: ModelDownloadProgress) => void,
+  ): Promise<void> {
+    const dir = await this.getDir();
+    await this.delete(filename);
+
+    const tempFilename = `${filename}${this.tempSuffix}`;
+    const tempHandle = await dir.getFileHandle(tempFilename, { create: true });
+    const tempWritable = await tempHandle.createWritable();
+    const reader = stream.getReader();
+    let receivedBytes = 0;
+
+    try {
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        if (!value) continue;
+        await tempWritable.write(value);
+        receivedBytes += value.byteLength;
+        onProgress?.({
+          receivedBytes,
+          totalBytes,
+          percent: totalBytes ? (receivedBytes / totalBytes) * 100 : 0,
+        });
+      }
+      await tempWritable.close();
+
+      const sizeValue = totalBytes ?? receivedBytes;
+      if (!Number.isFinite(sizeValue) || sizeValue <= 0) {
+        throw new Error('Model download did not provide a valid size.');
+      }
+
+      console.log(`[opfsModelCache] Finalizing cached model ${filename} from temp file ${tempFilename}.`);
+      const finalHandle = await dir.getFileHandle(filename, { create: true });
+      const finalWritable = await finalHandle.createWritable();
+      try {
+        const tempFile = await tempHandle.getFile();
+        console.log('[opfsModelCache] Temp file ready for finalize copy.', {
+          filename,
+          tempFilename,
+          tempSize: tempFile.size,
+          expectedSize: sizeValue,
+        });
+        await finalWritable.write(tempFile);
+        await finalWritable.close();
+      } catch (error) {
+        await finalWritable.abort();
+        throw error;
+      }
+
+      const sizeHandle = await dir.getFileHandle(this.getSizeFilename(filename), { create: true });
+      const sizeWritable = await sizeHandle.createWritable();
+      try {
+        await sizeWritable.write(String(sizeValue));
+        await sizeWritable.close();
+      } catch (error) {
+        await sizeWritable.abort();
+        throw error;
+      }
+
+      onProgress?.({
+        receivedBytes: sizeValue,
+        totalBytes: sizeValue,
+        percent: 100,
+      });
+      console.log('[opfsModelCache] Cached model finalize completed.', {
+        filename,
+        size: sizeValue,
+      });
+    } catch (error) {
+      try {
+        await tempWritable.abort();
+      } catch {
+        // Ignore abort cleanup errors.
+      }
+      await this.delete(filename);
+      throw error;
+    } finally {
+      await this.removeIfExists(dir, tempFilename);
+      reader.releaseLock();
+    }
+  }
+
+  private getSizeFilename(filename: string): string {
+    return `${filename}${this.sizeSuffix}`;
+  }
+
+  private async getDir(): Promise<FileSystemDirectoryHandle> {
+    const root = await navigator.storage.getDirectory();
+    return root.getDirectoryHandle(this.directoryName, { create: true });
+  }
+
+  private async removeIfExists(dir: FileSystemDirectoryHandle, name: string): Promise<void> {
+    try {
+      await dir.removeEntry(name);
+    } catch {
+      // Ignore missing entry cleanup.
+    }
+  }
+}