@lobehub/chat 1.36.46 → 1.37.0

package/src/app/(main)/settings/_layout/Desktop/index.tsx

@@ -6,6 +6,7 @@ import { memo, useRef } from 'react';
 import { useTranslation } from 'react-i18next';
 import { Flexbox } from 'react-layout-kit';
 
+import InitClientDB from '@/features/InitClientDB';
 import Footer from '@/features/Setting/Footer';
 import SettingContainer from '@/features/Setting/SettingContainer';
 import { useActiveSettingsKey } from '@/hooks/useActiveSettingsKey';
@@ -45,6 +46,7 @@ const Layout = memo<LayoutProps>(({ children, category }) => {
         </Header>
       )}
       <SettingContainer addonAfter={<Footer />}>{children}</SettingContainer>
+      <InitClientDB />
     </Flexbox>
   );
 });

package/src/app/loading/Client/Content.tsx

@@ -1,17 +1,25 @@
+import dynamic from 'next/dynamic';
 import React, { memo } from 'react';
 import { useTranslation } from 'react-i18next';
 
 import FullscreenLoading from '@/components/FullscreenLoading';
 import { useGlobalStore } from '@/store/global';
 import { systemStatusSelectors } from '@/store/global/selectors';
+import { DatabaseLoadingState } from '@/types/clientDB';
 
 import { CLIENT_LOADING_STAGES } from '../stage';
 
+const InitError = dynamic(() => import('./Error'), { ssr: false });
+
 interface InitProps {
   setActiveStage: (value: string) => void;
 }
 
-const Init = memo<InitProps>(() => {
+const Init = memo<InitProps>(({ setActiveStage }) => {
+  const useInitClientDB = useGlobalStore((s) => s.useInitClientDB);
+
+  useInitClientDB({ onStateChange: setActiveStage });
+
   return null;
 });
 
@@ -23,12 +31,14 @@ interface ContentProps {
 const Content = memo<ContentProps>(({ loadingStage, setActiveStage }) => {
   const { t } = useTranslation('common');
   const isPgliteNotInited = useGlobalStore(systemStatusSelectors.isPgliteNotInited);
+  const isError = useGlobalStore((s) => s.initClientDBStage === DatabaseLoadingState.Error);
 
   return (
     <>
       {isPgliteNotInited && <Init setActiveStage={setActiveStage} />}
       <FullscreenLoading
         activeStage={CLIENT_LOADING_STAGES.indexOf(loadingStage)}
+        contentRender={isError && <InitError />}
         stages={CLIENT_LOADING_STAGES.map((key) => t(`appLoading.${key}` as any))}
       />
     </>

package/src/app/loading/Client/Error.tsx

@@ -0,0 +1,27 @@
+import { Button } from 'antd';
+import React from 'react';
+import { useTranslation } from 'react-i18next';
+import { Center } from 'react-layout-kit';
+
+import ErrorResult from '@/features/InitClientDB/ErrorResult';
+
+const InitError = () => {
+  const { t } = useTranslation('common');
+
+  return (
+    <ErrorResult>
+      {({ setOpen }) => (
+        <Center gap={8}>
+          {t('appLoading.failed')}
+          <div>
+            <Button onClick={() => setOpen(true)} type={'primary'}>
+              {t('appLoading.showDetail')}
+            </Button>
+          </div>
+        </Center>
+      )}
+    </ErrorResult>
+  );
+};
+
+export default InitError;

package/src/app/loading/stage.ts

@@ -1,3 +1,5 @@
+import { DatabaseLoadingState } from '@/types/clientDB';
+
 export enum AppLoadingStage {
   GoToChat = 'goToChat',
   Idle = 'appIdle',
@@ -17,6 +19,12 @@ export const SERVER_LOADING_STAGES = [
 export const CLIENT_LOADING_STAGES = [
   AppLoadingStage.Idle,
   AppLoadingStage.Initializing,
+  DatabaseLoadingState.Initializing,
+  DatabaseLoadingState.LoadingDependencies,
+  DatabaseLoadingState.LoadingWasm,
+  DatabaseLoadingState.Migrating,
+  DatabaseLoadingState.Finished,
+  DatabaseLoadingState.Ready,
   AppLoadingStage.InitUser,
   AppLoadingStage.GoToChat,
 ] as string[];

package/src/components/FullscreenLoading/index.tsx

@@ -1,4 +1,4 @@
-import React, { memo } from 'react';
+import React, { ReactNode, memo } from 'react';
 import { Center, Flexbox } from 'react-layout-kit';
 
 import { ProductLogo } from '@/components/Branding';
@@ -6,15 +6,16 @@ import InitProgress, { StageItem } from '@/components/InitProgress';
 
 interface FullscreenLoadingProps {
   activeStage: number;
+  contentRender?: ReactNode;
   stages: StageItem[];
 }
 
-const FullscreenLoading = memo<FullscreenLoadingProps>(({ activeStage, stages }) => {
+const FullscreenLoading = memo<FullscreenLoadingProps>(({ activeStage, stages, contentRender }) => {
   return (
     <Flexbox height={'100%'} style={{ position: 'relative', userSelect: 'none' }} width={'100%'}>
       <Center flex={1} gap={16} width={'100%'}>
         <ProductLogo size={48} type={'combine'} />
-        <InitProgress activeStage={activeStage} stages={stages} />
+        {contentRender ? contentRender : <InitProgress activeStage={activeStage} stages={stages} />}
       </Center>
     </Flexbox>
   );

package/src/const/version.ts

@@ -6,6 +6,7 @@ import { BRANDING_NAME, ORG_NAME } from './branding';
 export const CURRENT_VERSION = pkg.version;
 
 export const isServerMode = getServerDBConfig().NEXT_PUBLIC_ENABLED_SERVER_SERVICE;
+export const isUsePgliteDB = process.env.NEXT_PUBLIC_CLIENT_DB === 'pglite';
 
 // @ts-ignore
 export const isCustomBranding = BRANDING_NAME !== 'LobeChat';

package/src/database/client/db.test.ts

@@ -0,0 +1,172 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { ClientDBLoadingProgress, DatabaseLoadingState } from '@/types/clientDB';
+
+import { DatabaseManager } from './db';
+
+// Mock 所有外部依赖
+vi.mock('@electric-sql/pglite', () => ({
+  default: vi.fn(),
+  IdbFs: vi.fn(),
+  PGlite: vi.fn(),
+  MemoryFS: vi.fn(),
+}));
+
+vi.mock('@electric-sql/pglite/vector', () => ({
+  default: vi.fn(),
+  vector: vi.fn(),
+}));
+
+vi.mock('drizzle-orm/pglite', () => ({
+  drizzle: vi.fn(() => ({
+    dialect: {
+      migrate: vi.fn().mockResolvedValue(undefined),
+    },
+  })),
+}));
+
+let manager: DatabaseManager;
+let progressEvents: ClientDBLoadingProgress[] = [];
+let stateChanges: DatabaseLoadingState[] = [];
+
+let callbacks = {
+  onProgress: vi.fn((progress: ClientDBLoadingProgress) => {
+    progressEvents.push(progress);
+  }),
+  onStateChange: vi.fn((state: DatabaseLoadingState) => {
+    stateChanges.push(state);
+  }),
+};
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  progressEvents = [];
+  stateChanges = [];
+
+  callbacks = {
+    onProgress: vi.fn((progress: ClientDBLoadingProgress) => {
+      progressEvents.push(progress);
+    }),
+    onStateChange: vi.fn((state: DatabaseLoadingState) => {
+      stateChanges.push(state);
+    }),
+  };
+  // @ts-expect-error
+  DatabaseManager['instance'] = undefined;
+  manager = DatabaseManager.getInstance();
+});
+
+describe('DatabaseManager', () => {
+  describe('Callback Handling', () => {
+    it('should properly track loading states', async () => {
+      await manager.initialize(callbacks);
+
+      // 验证状态转换顺序
+      expect(stateChanges).toEqual([
+        DatabaseLoadingState.Initializing,
+        DatabaseLoadingState.LoadingDependencies,
+        DatabaseLoadingState.LoadingWasm,
+        DatabaseLoadingState.Migrating,
+        DatabaseLoadingState.Finished,
+        DatabaseLoadingState.Ready,
+      ]);
+    });
+
+    it('should report dependencies loading progress', async () => {
+      await manager.initialize(callbacks);
+
+      // 验证依赖加载进度回调
+      const dependencyProgress = progressEvents.filter((e) => e.phase === 'dependencies');
+      expect(dependencyProgress.length).toBeGreaterThan(0);
+      expect(dependencyProgress[dependencyProgress.length - 1]).toEqual(
+        expect.objectContaining({
+          phase: 'dependencies',
+          progress: 100,
+          costTime: expect.any(Number),
+        }),
+      );
+    });
+
+    it('should report WASM loading progress', async () => {
+      await manager.initialize(callbacks);
+
+      // 验证 WASM 加载进度回调
+      const wasmProgress = progressEvents.filter((e) => e.phase === 'wasm');
+      // expect(wasmProgress.length).toBeGreaterThan(0);
+      expect(wasmProgress[wasmProgress.length - 1]).toEqual(
+        expect.objectContaining({
+          phase: 'wasm',
+          progress: 100,
+          costTime: expect.any(Number),
+        }),
+      );
+    });
+
+    it('should handle initialization errors', async () => {
+      // 模拟加载失败
+      vi.spyOn(global, 'fetch').mockRejectedValueOnce(new Error('Network error'));
+
+      await expect(manager.initialize(callbacks)).rejects.toThrow();
+      expect(stateChanges).toContain(DatabaseLoadingState.Error);
+    });
+
+    it('should only initialize once when called multiple times', async () => {
+      const firstInit = manager.initialize(callbacks);
+      const secondInit = manager.initialize(callbacks);
+
+      await Promise.all([firstInit, secondInit]);
+
+      // 验证回调只被调用一次
+      const readyStateCount = stateChanges.filter(
+        (state) => state === DatabaseLoadingState.Ready,
+      ).length;
+      expect(readyStateCount).toBe(1);
+    });
+  });
+
+  describe('Progress Calculation', () => {
+    it('should report progress between 0 and 100', async () => {
+      await manager.initialize(callbacks);
+
+      // 验证所有进度值都在有效范围内
+      progressEvents.forEach((event) => {
+        expect(event.progress).toBeGreaterThanOrEqual(0);
+        expect(event.progress).toBeLessThanOrEqual(100);
+      });
+    });
+
+    it('should include timing information', async () => {
+      await manager.initialize(callbacks);
+
+      // 验证最终进度回调包含耗时信息
+      const finalProgress = progressEvents[progressEvents.length - 1];
+      expect(finalProgress.costTime).toBeGreaterThan(0);
+    });
+  });
+
+  describe('Error Handling', () => {
+    it('should handle missing callbacks gracefully', async () => {
+      // 测试没有提供回调的情况
+      await expect(manager.initialize()).resolves.toBeDefined();
+    });
+
+    it('should handle partial callbacks', async () => {
+      // 只提供部分回调
+      await expect(manager.initialize({ onProgress: callbacks.onProgress })).resolves.toBeDefined();
+      await expect(
+        manager.initialize({ onStateChange: callbacks.onStateChange }),
+      ).resolves.toBeDefined();
+    });
+  });
+
+  describe('Database Access', () => {
+    it('should throw error when accessing database before initialization', () => {
+      expect(() => manager.db).toThrow('Database not initialized');
+    });
+
+    it('should provide access to database after initialization', async () => {
+      await manager.initialize();
+      expect(() => manager.db).not.toThrow();
+    });
+  });
+});

package/src/database/client/db.ts

@@ -0,0 +1,246 @@
+import type { PgliteDatabase } from 'drizzle-orm/pglite';
+import { Md5 } from 'ts-md5';
+
+import { ClientDBLoadingProgress, DatabaseLoadingState } from '@/types/clientDB';
+import { sleep } from '@/utils/sleep';
+
+import * as schema from '../schemas';
+import migrations from './migrations.json';
+
+const pgliteSchemaHashCache = 'LOBE_CHAT_PGLITE_SCHEMA_HASH';
+
+type DrizzleInstance = PgliteDatabase<typeof schema>;
+
+export interface DatabaseLoadingCallbacks {
+  onError?: (error: Error) => void;
+  onProgress?: (progress: ClientDBLoadingProgress) => void;
+  onStateChange?: (state: DatabaseLoadingState) => void;
+}
+
+export class DatabaseManager {
+  private static instance: DatabaseManager;
+  private dbInstance: DrizzleInstance | null = null;
+  private initPromise: Promise<DrizzleInstance> | null = null;
+  private callbacks?: DatabaseLoadingCallbacks;
+  private isLocalDBSchemaSynced = false;
+
+  // CDN 配置
+  private static WASM_CDN_URL =
+    'https://registry.npmmirror.com/@electric-sql/pglite/0.2.13/files/dist/postgres.wasm';
+
+  private constructor() {}
+
+  static getInstance() {
+    if (!DatabaseManager.instance) {
+      DatabaseManager.instance = new DatabaseManager();
+    }
+    return DatabaseManager.instance;
+  }
+
+  // 加载并编译 WASM 模块
+  private async loadWasmModule(): Promise<WebAssembly.Module> {
+    const start = Date.now();
+    this.callbacks?.onStateChange?.(DatabaseLoadingState.LoadingWasm);
+
+    const response = await fetch(DatabaseManager.WASM_CDN_URL);
+
+    const contentLength = Number(response.headers.get('Content-Length')) || 0;
+    const reader = response.body?.getReader();
+
+    if (!reader) throw new Error('Failed to start WASM download');
+
+    let receivedLength = 0;
+    const chunks: Uint8Array[] = [];
+
+    // 读取数据流
+    // eslint-disable-next-line no-constant-condition
+    while (true) {
+      const { done, value } = await reader.read();
+
+      if (done) break;
+
+      chunks.push(value);
+      receivedLength += value.length;
+
+      // 计算并报告进度
+      const progress = Math.min(Math.round((receivedLength / contentLength) * 100), 100);
+      this.callbacks?.onProgress?.({
+        phase: 'wasm',
+        progress,
+      });
+    }
+
+    // 合并数据块
+    const wasmBytes = new Uint8Array(receivedLength);
+    let position = 0;
+    for (const chunk of chunks) {
+      wasmBytes.set(chunk, position);
+      position += chunk.length;
+    }
+
+    this.callbacks?.onProgress?.({
+      costTime: Date.now() - start,
+      phase: 'wasm',
+      progress: 100,
+    });
+
+    // 编译 WASM 模块
+    return WebAssembly.compile(wasmBytes);
+  }
+
+  // 异步加载 PGlite 相关依赖
+  private async loadDependencies() {
+    const start = Date.now();
+    this.callbacks?.onStateChange?.(DatabaseLoadingState.LoadingDependencies);
+
+    const imports = [
+      import('@electric-sql/pglite').then((m) => ({
+        IdbFs: m.IdbFs,
+        MemoryFS: m.MemoryFS,
+        PGlite: m.PGlite,
+      })),
+      import('@electric-sql/pglite/vector'),
+      import('drizzle-orm/pglite'),
+    ];
+
+    let loaded = 0;
+    const results = await Promise.all(
+      imports.map(async (importPromise) => {
+        const result = await importPromise;
+        loaded += 1;
+
+        // 计算加载进度
+        this.callbacks?.onProgress?.({
+          phase: 'dependencies',
+          progress: Math.min(Math.round((loaded / imports.length) * 100), 100),
+        });
+        return result;
+      }),
+    );
+
+    this.callbacks?.onProgress?.({
+      costTime: Date.now() - start,
+      phase: 'dependencies',
+      progress: 100,
+    });
+
+    // @ts-ignore
+    const [{ PGlite, IdbFs, MemoryFS }, { vector }, { drizzle }] = results;
+
+    return { IdbFs, MemoryFS, PGlite, drizzle, vector };
+  }
+
+  // 数据库迁移方法
+  private async migrate(skipMultiRun = false): Promise<DrizzleInstance> {
+    if (this.isLocalDBSchemaSynced && skipMultiRun) return this.db;
+
+    const cacheHash = localStorage.getItem(pgliteSchemaHashCache);
+    const hash = Md5.hashStr(JSON.stringify(migrations));
+
+    // if hash is the same, no need to migrate
+    if (hash === cacheHash) {
+      this.isLocalDBSchemaSynced = true;
+      return this.db;
+    }
+
+    const start = Date.now();
+    try {
+      this.callbacks?.onStateChange?.(DatabaseLoadingState.Migrating);
+
+      // refs: https://github.com/drizzle-team/drizzle-orm/discussions/2532
+      // @ts-expect-error
+      await this.db.dialect.migrate(migrations, this.db.session, {});
+      localStorage.setItem(pgliteSchemaHashCache, hash);
+      this.isLocalDBSchemaSynced = true;
+
+      console.info(`🗂 Migration success, take ${Date.now() - start}ms`);
+    } catch (cause) {
+      console.error('❌ Local database schema migration failed', cause);
+      throw cause;
+    }
+
+    return this.db;
+  }
+
+  // 初始化数据库
+  async initialize(callbacks?: DatabaseLoadingCallbacks): Promise<DrizzleInstance> {
+    if (this.initPromise) return this.initPromise;
+
+    this.callbacks = callbacks;
+
+    this.initPromise = (async () => {
+      try {
+        if (this.dbInstance) return this.dbInstance;
+
+        const time = Date.now();
+        // 初始化数据库
+        this.callbacks?.onStateChange?.(DatabaseLoadingState.Initializing);
+
+        // 加载依赖
+        const { PGlite, vector, drizzle, IdbFs, MemoryFS } = await this.loadDependencies();
+
+        // 加载并编译 WASM 模块
+        const wasmModule = await this.loadWasmModule();
+
+        const db = new PGlite({
+          extensions: { vector },
+          fs: typeof window === 'undefined' ? new MemoryFS('lobechat') : new IdbFs('lobechat'),
+          relaxedDurability: true,
+          wasmModule,
+        });
+
+        this.dbInstance = drizzle({ client: db, schema });
+
+        await this.migrate(true);
+
+        this.callbacks?.onStateChange?.(DatabaseLoadingState.Finished);
+        console.log(`✅ Database initialized in ${Date.now() - time}ms`);
+
+        await sleep(50);
+
+        this.callbacks?.onStateChange?.(DatabaseLoadingState.Ready);
+
+        return this.dbInstance as DrizzleInstance;
+      } catch (e) {
+        this.initPromise = null;
+        this.callbacks?.onStateChange?.(DatabaseLoadingState.Error);
+        const error = e as Error;
+        this.callbacks?.onError?.({
+          message: error.message,
+          name: error.name,
+          stack: error.stack,
+        });
+        throw error;
+      }
+    })();
+
+    return this.initPromise;
+  }
+
+  // 获取数据库实例
+  get db(): DrizzleInstance {
+    if (!this.dbInstance) {
+      throw new Error('Database not initialized. Please call initialize() first.');
+    }
+    return this.dbInstance;
+  }
+
+  // 创建代理对象
+  createProxy(): DrizzleInstance {
+    return new Proxy({} as DrizzleInstance, {
+      get: (target, prop) => {
+        return this.db[prop as keyof DrizzleInstance];
+      },
+    });
+  }
+}
+
+// 导出单例
+const dbManager = DatabaseManager.getInstance();
+
+// 保持原有的 clientDB 导出不变
+export const clientDB = dbManager.createProxy();
+
+// 导出初始化方法，供应用启动时使用
+export const initializeDB = (callbacks?: DatabaseLoadingCallbacks) =>
+  dbManager.initialize(callbacks);