@lobehub/lobehub 2.0.0-next.190 → 2.0.0-next.192

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

@@ -1,18 +1,11 @@
-import { ClientDBLoadingProgress, DatabaseLoadingState } from '@lobechat/types';
+import { PGlite } from '@electric-sql/pglite';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 
-import { DatabaseManager } from './db';
-
-// Mock 所有外部依赖
 vi.mock('@electric-sql/pglite', () => ({
-  default: vi.fn(),
-  IdbFs: vi.fn(),
-  PGlite: vi.fn(),
-  MemoryFS: vi.fn(),
+  PGlite: vi.fn(() => ({})),
 }));
 
 vi.mock('@electric-sql/pglite/vector', () => ({
-  default: vi.fn(),
   vector: vi.fn(),
 }));
 
@@ -24,154 +17,36 @@ vi.mock('drizzle-orm/pglite', () => ({
   })),
 }));
 
-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();
+  vi.resetModules();
 });
 
 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,
-        ]);
-      },
-      {
-        timeout: 15000,
-      },
-    );
-
-    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);
+  describe('initializeDB', () => {
+    it('should initialize database with PGlite', async () => {
+      const { initializeDB } = await import('./db');
+      await initializeDB();
+
+      expect(PGlite).toHaveBeenCalledWith('idb://lobechat', {
+        extensions: { vector: expect.any(Function) },
+        relaxedDurability: true,
       });
     });
 
-    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 only initialize once when called multiple times', async () => {
+      const { initializeDB } = await import('./db');
+      await Promise.all([initializeDB(), initializeDB()]);
 
-    it('should handle partial callbacks', async () => {
-      // 只提供部分回调
-      await expect(manager.initialize({ onProgress: callbacks.onProgress })).resolves.toBeDefined();
-      await expect(
-        manager.initialize({ onStateChange: callbacks.onStateChange }),
-      ).resolves.toBeDefined();
+      expect(PGlite).toHaveBeenCalledTimes(1);
     });
   });
 
-  describe('Database Access', () => {
-    it('should throw error when accessing database before initialization', () => {
-      expect(() => manager.db).toThrow('Database not initialized');
-    });
-
+  describe('clientDB proxy', () => {
     it('should provide access to database after initialization', async () => {
-      await manager.initialize();
-      expect(() => manager.db).not.toThrow();
+      const { clientDB, initializeDB } = await import('./db');
+      await initializeDB();
+      expect(clientDB).toBeDefined();
     });
   });
 });

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

@@ -1,53 +1,24 @@
-import {
-  type ClientDBLoadingProgress,
-  DatabaseLoadingState,
-  type MigrationSQL,
-  type MigrationTableItem,
-} from '@lobechat/types';
+import { PGlite } from '@electric-sql/pglite';
+import { vector } from '@electric-sql/pglite/vector';
 import { sql } from 'drizzle-orm';
 import { PgliteDatabase, drizzle } from 'drizzle-orm/pglite';
 import { Md5 } from 'ts-md5';
 
-import { sleep } from '@/utils/sleep';
-
 import migrations from '../core/migrations.json';
 import { DrizzleMigrationModel } from '../models/drizzleMigration';
 import * as schema from '../schemas';
 
 const pgliteSchemaHashCache = 'LOBE_CHAT_PGLITE_SCHEMA_HASH';
-
 const DB_NAME = 'lobechat';
-type DrizzleInstance = PgliteDatabase<typeof schema>;
 
-interface onErrorState {
-  error: Error;
-  migrationTableItems: MigrationTableItem[];
-  migrationsSQL: MigrationSQL[];
-}
-
-export interface DatabaseLoadingCallbacks {
-  onError?: (error: onErrorState) => void;
-  onProgress?: (progress: ClientDBLoadingProgress) => void;
-  onStateChange?: (state: DatabaseLoadingState) => void;
-}
+type DrizzleInstance = PgliteDatabase<typeof schema>;
 
-export class DatabaseManager {
+class DatabaseManager {
   private static instance: DatabaseManager;
   private dbInstance: DrizzleInstance | null = null;
   private initPromise: Promise<DrizzleInstance> | null = null;
-  private callbacks?: DatabaseLoadingCallbacks;
   private isLocalDBSchemaSynced = false;
 
-  // CDN configuration
-  private static WASM_CDN_URL =
-    'https://registry.npmmirror.com/@electric-sql/pglite/0.2.17/files/dist/postgres.wasm';
-
-  private static FSBUNDLER_CDN_URL =
-    'https://registry.npmmirror.com/@electric-sql/pglite/0.2.17/files/dist/postgres.data';
-
-  private static VECTOR_CDN_URL =
-    'https://registry.npmmirror.com/@electric-sql/pglite/0.2.17/files/dist/vector.tar.gz';
-
   private constructor() {}
 
   static getInstance() {
@@ -57,108 +28,8 @@ export class DatabaseManager {
     return DatabaseManager.instance;
   }
 
-  // Load and compile WASM module
-  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[] = [];
-
-    // Read data stream
-    // eslint-disable-next-line no-constant-condition
-    while (true) {
-      const { done, value } = await reader.read();
-
-      if (done) break;
-
-      chunks.push(value);
-      receivedLength += value.length;
-
-      // Calculate and report progress
-      const progress = Math.min(Math.round((receivedLength / contentLength) * 100), 100);
-      this.callbacks?.onProgress?.({
-        phase: 'wasm',
-        progress,
-      });
-    }
-
-    // Merge data chunks
-    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,
-    });
-
-    // Compile WASM module
-    return WebAssembly.compile(wasmBytes);
-  }
-
-  private fetchFsBundle = async () => {
-    const res = await fetch(DatabaseManager.FSBUNDLER_CDN_URL);
-
-    return await res.blob();
-  };
-
-  // Asynchronously load PGlite related dependencies
-  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'),
-      this.fetchFsBundle(),
-    ];
-
-    let loaded = 0;
-    const results = await Promise.all(
-      imports.map(async (importPromise) => {
-        const result = await importPromise;
-        loaded += 1;
-
-        // Calculate loading progress
-        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 }, fsBundle] = results;
-
-    return { IdbFs, MemoryFS, PGlite, fsBundle, vector };
-  }
-
-  // Database migration method
-  private async migrate(skipMultiRun = false): Promise<DrizzleInstance> {
-    if (this.isLocalDBSchemaSynced && skipMultiRun) return this.db;
+  private async migrate(): Promise<DrizzleInstance> {
+    if (this.isLocalDBSchemaSynced) return this.db;
 
     let hash: string | undefined;
     if (typeof localStorage !== 'undefined') {
@@ -179,17 +50,13 @@ export class DatabaseManager {
           }
         } catch (error) {
           console.warn('Error checking table existence, proceeding with migration', error);
-          // If query fails, continue migration to ensure safety
         }
       }
     }
 
     const start = Date.now();
     try {
-      this.callbacks?.onStateChange?.(DatabaseLoadingState.Migrating);
-
-      // refs: https://github.com/drizzle-team/drizzle-orm/discussions/2532
-      // @ts-expect-error
+      // @ts-expect-error - migrate internal API
       await this.db.dialect.migrate(migrations, this.db.session, {});
 
       if (typeof localStorage !== 'undefined' && hash) {
@@ -197,7 +64,6 @@ export class DatabaseManager {
       }
 
       this.isLocalDBSchemaSynced = true;
-
       console.info(`🗂 Migration success, take ${Date.now() - start}ms`);
     } catch (cause) {
       console.error('❌ Local database schema migration failed', cause);
@@ -207,95 +73,32 @@ export class DatabaseManager {
     return this.db;
   }
 
-  // Initialize database
-  async initialize(callbacks?: DatabaseLoadingCallbacks): Promise<DrizzleInstance> {
+  async initialize(): 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();
-        // Initialize database
-        this.callbacks?.onStateChange?.(DatabaseLoadingState.Initializing);
-
-        // Load dependencies
-        const { fsBundle, PGlite, MemoryFS, IdbFs, vector } = await this.loadDependencies();
-
-        // Load and compile WASM module
-        const wasmModule = await this.loadWasmModule();
-
-        const { initPgliteWorker } = await import('./pglite');
-
-        let db: typeof PGlite;
-
-        // make db as web worker if worker is available
-        // https://github.com/lobehub/lobe-chat/issues/5785
-        if (typeof Worker !== 'undefined' && typeof navigator.locks !== 'undefined') {
-          db = await initPgliteWorker({
-            dbName: DB_NAME,
-            fsBundle: fsBundle as Blob,
-            vectorBundlePath: DatabaseManager.VECTOR_CDN_URL,
-            wasmModule,
-          });
-        } else {
-          // in edge runtime or test runtime, we don't have worker
-          db = new PGlite({
-            extensions: { vector },
-            fs: typeof window === 'undefined' ? new MemoryFS(DB_NAME) : new IdbFs(DB_NAME),
-            relaxedDurability: true,
-            wasmModule,
-          });
-        }
-
-        this.dbInstance = drizzle({ client: db, schema });
+      if (this.dbInstance) return this.dbInstance;
 
-        await this.migrate(true);
+      const time = Date.now();
 
-        this.callbacks?.onStateChange?.(DatabaseLoadingState.Finished);
-        console.log(`✅ Database initialized in ${Date.now() - time}ms`);
-
-        await sleep(50);
+      // 直接使用 pglite，自动处理 wasm 加载
+      const pglite = new PGlite(`idb://${DB_NAME}`, {
+        extensions: { vector },
+        relaxedDurability: true,
+      });
 
-        this.callbacks?.onStateChange?.(DatabaseLoadingState.Ready);
+      this.dbInstance = drizzle({ client: pglite, schema });
 
-        return this.dbInstance as DrizzleInstance;
-      } catch (e) {
-        this.initPromise = null;
-        this.callbacks?.onStateChange?.(DatabaseLoadingState.Error);
-        const error = e as Error;
+      await this.migrate();
 
-        // Query migration table data
-        let migrationsTableData: MigrationTableItem[] = [];
-        try {
-          // Attempt to query migration table
-          const drizzleMigration = new DrizzleMigrationModel(this.db as any);
-          migrationsTableData = await drizzleMigration.getMigrationList();
-        } catch (queryError) {
-          console.error('Failed to query migrations table:', queryError);
-        }
+      console.log(`✅ Database initialized in ${Date.now() - time}ms`);
 
-        this.callbacks?.onError?.({
-          error: {
-            message: error.message,
-            name: error.name,
-            stack: error.stack,
-          },
-          migrationTableItems: migrationsTableData,
-          migrationsSQL: migrations,
-        });
-
-        console.error(error);
-        throw error;
-      }
+      return this.dbInstance;
     })();
 
     return this.initPromise;
   }
 
-  // Get database instance
   get db(): DrizzleInstance {
     if (!this.dbInstance) {
       throw new Error('Database not initialized. Please call initialize() first.');
@@ -303,7 +106,6 @@ export class DatabaseManager {
     return this.dbInstance;
   }
 
-  // Create proxy object
   createProxy(): DrizzleInstance {
     return new Proxy({} as DrizzleInstance, {
       get: (target, prop) => {
@@ -313,7 +115,7 @@ export class DatabaseManager {
   }
 
   async resetDatabase(): Promise<void> {
-    // 1. Close existing PGlite connection (if exists)
+    // 1. Close existing PGlite connection
     if (this.dbInstance) {
       try {
         // @ts-ignore
@@ -321,31 +123,28 @@ export class DatabaseManager {
         console.log('PGlite instance closed successfully.');
       } catch (e) {
         console.error('Error closing PGlite instance:', e);
-        // Even if closing fails, continue with deletion attempt; IndexedDB onblocked or onerror will handle subsequent issues
       }
     }
 
     // 2. Reset database instance and initialization state
     this.dbInstance = null;
     this.initPromise = null;
-    this.isLocalDBSchemaSynced = false; // Reset sync state
+    this.isLocalDBSchemaSynced = false;
 
     // 3. Delete IndexedDB database
     return new Promise<void>((resolve, reject) => {
-      // Check if IndexedDB is available
       if (typeof indexedDB === 'undefined') {
         console.warn('IndexedDB is not available, cannot delete database');
-        resolve(); // Cannot delete in this environment, resolve directly
+        resolve();
         return;
       }
 
-      const dbName = `/pglite/${DB_NAME}`; // Path used by PGlite IdbFs
+      const dbName = `/pglite/${DB_NAME}`;
       const request = indexedDB.deleteDatabase(dbName);
 
       request.onsuccess = () => {
         console.log(`✅ Database '${dbName}' reset successfully`);
 
-        // Clear locally stored schema hash
         if (typeof localStorage !== 'undefined') {
           localStorage.removeItem(pgliteSchemaHashCache);
         }
@@ -365,14 +164,10 @@ export class DatabaseManager {
       };
 
       request.onblocked = (event) => {
-        // This event is triggered when other open connections block database deletion
-        console.warn(
-          `Deletion of database '${dbName}' is blocked. This usually means other connections (e.g., in other tabs) are still open. Event:`,
-          event,
-        );
+        console.warn(`Deletion of database '${dbName}' is blocked.`, event);
         reject(
           new Error(
-            `Failed to reset database '${dbName}' because it is blocked by other open connections. Please close other tabs or applications using this database and try again.`,
+            `Failed to reset database '${dbName}' because it is blocked by other open connections.`,
           ),
         );
       };
@@ -383,12 +178,9 @@ export class DatabaseManager {
 // Export singleton
 const dbManager = DatabaseManager.getInstance();
 
-// Keep original clientDB export unchanged
 export const clientDB = dbManager.createProxy();
 
-// Export initialization method for application startup
-export const initializeDB = (callbacks?: DatabaseLoadingCallbacks) =>
-  dbManager.initialize(callbacks);
+export const initializeDB = () => dbManager.initialize();
 
 export const resetClientDatabase = async () => {
   await dbManager.resetDatabase();

package/packages/database/src/models/__tests__/_util.ts

@@ -1,14 +1,30 @@
-import { clientDB, initializeDB } from '../../client/db';
+import { PGlite } from '@electric-sql/pglite';
+import { vector } from '@electric-sql/pglite/vector';
+import { drizzle } from 'drizzle-orm/pglite';
+
+import migrations from '../../core/migrations.json';
+import * as schema from '../../schemas';
 import { LobeChatDatabase } from '../../type';
 
 const isServerDBMode = process.env.TEST_SERVER_DB === '1';
 
+let testClientDB: ReturnType<typeof drizzle<typeof schema>> | null = null;
+
 export const getTestDB = async () => {
   if (isServerDBMode) {
     const { getTestDBInstance } = await import('../../core/dbForTest');
     return await getTestDBInstance();
   }
 
-  await initializeDB();
-  return clientDB as LobeChatDatabase;
+  if (testClientDB) return testClientDB as unknown as LobeChatDatabase;
+
+  // 直接使用 pglite 内置资源，不需要从 CDN 下载
+  const pglite = new PGlite({ extensions: { vector } });
+
+  testClientDB = drizzle({ client: pglite, schema });
+
+  // @ts-expect-error - migrate internal API
+  await testClientDB.dialect.migrate(migrations, testClientDB.session, {});
+
+  return testClientDB as unknown as LobeChatDatabase;
 };

package/packages/database/src/models/message.ts

@@ -485,7 +485,6 @@ export class MessageModel {
             .filter((relation) => relation.messageId === item.id)
             // eslint-disable-next-line @typescript-eslint/no-unused-vars
             .map<ChatImageItem>(({ id, url, name }) => ({ alt: name!, id, url })),
-          meta: {},
 
           model,
 

package/packages/prompts/src/prompts/chatMessages/index.test.ts

@@ -15,7 +15,6 @@ const createMessage = (
   id: overrides.id ?? `msg-${++messageCounter}`,
   createdAt: overrides.createdAt ?? 0,
   updatedAt: overrides.updatedAt ?? 0,
-  meta: overrides.meta ?? {},
   ...overrides,
 });