@lobehub/lobehub 2.0.0-next.191 → 2.0.0-next.193

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/__tests__/knowledgeBase.test.ts

@@ -2,7 +2,7 @@
 import { and, eq } from 'drizzle-orm';
 import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
-import { LobeChatDatabase } from '../../type';import { sleep } from '@/utils/sleep';
+import { sleep } from '@/utils/sleep';
 
 import {
   NewKnowledgeBase,
@@ -12,6 +12,7 @@ import {
   knowledgeBases,
   users,
 } from '../../schemas';
+import { LobeChatDatabase } from '../../type';
 import { KnowledgeBaseModel } from '../knowledgeBase';
 import { getTestDB } from './_util';
 
@@ -228,6 +229,34 @@ describe('KnowledgeBaseModel', () => {
       expect(remainingFiles).toHaveLength(1);
       expect(remainingFiles[0].fileId).toBe('file2');
     });
+
+    it('should not allow removing files from another user knowledge base', async () => {
+      await serverDB.insert(globalFiles).values([
+        {
+          hashId: 'hash1',
+          url: 'https://example.com/document.pdf',
+          size: 1000,
+          fileType: 'application/pdf',
+          creator: userId,
+        },
+      ]);
+
+      await serverDB.insert(files).values([fileList[0]]);
+
+      const { id: knowledgeBaseId } = await knowledgeBaseModel.create({ name: 'Test Group' });
+      await knowledgeBaseModel.addFilesToKnowledgeBase(knowledgeBaseId, ['file1']);
+
+      // Another user tries to remove files from this knowledge base
+      const attackerModel = new KnowledgeBaseModel(serverDB, 'user2');
+      await attackerModel.removeFilesFromKnowledgeBase(knowledgeBaseId, ['file1']);
+
+      // Files should still exist since the attacker doesn't own them
+      const remainingFiles = await serverDB.query.knowledgeBaseFiles.findMany({
+        where: eq(knowledgeBaseFiles.knowledgeBaseId, knowledgeBaseId),
+      });
+      expect(remainingFiles).toHaveLength(1);
+      expect(remainingFiles[0].fileId).toBe('file1');
+    });
   });
 
   describe('static findById', () => {

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

@@ -43,13 +43,15 @@ export class KnowledgeBaseModel {
   };
 
   removeFilesFromKnowledgeBase = async (knowledgeBaseId: string, ids: string[]) => {
-    return this.db.delete(knowledgeBaseFiles).where(
-      and(
-        eq(knowledgeBaseFiles.knowledgeBaseId, knowledgeBaseId),
-        inArray(knowledgeBaseFiles.fileId, ids),
-        // eq(knowledgeBaseFiles.userId, this.userId),
-      ),
-    );
+    return this.db
+      .delete(knowledgeBaseFiles)
+      .where(
+        and(
+          eq(knowledgeBaseFiles.userId, this.userId),
+          eq(knowledgeBaseFiles.knowledgeBaseId, knowledgeBaseId),
+          inArray(knowledgeBaseFiles.fileId, ids),
+        ),
+      );
   };
   // query
   query = async () => {

package/src/app/[variants]/(main)/settings/provider/features/ModelList/ModelItem.tsx

@@ -22,13 +22,14 @@ import ModelConfigModal from './ModelConfigModal';
 import { ProviderSettingsContext } from './ProviderSettingsContext';
 
 const styles = createStaticStyles(({ css, cx }) => {
-  const config = css`
-    opacity: 0;
-    transition: all 100ms ease-in-out;
-  `;
-
   return {
-    config,
+    config: cx(
+      'model-item-config',
+      css`
+        opacity: 0;
+        transition: all 100ms ease-in-out;
+      `,
+    ),
     container: css`
       position: relative;
       border-radius: ${cssVar.borderRadiusLG}px;
@@ -37,7 +38,7 @@ const styles = createStaticStyles(({ css, cx }) => {
       &:hover {
         background-color: ${cssVar.colorFillTertiary};
 
-        .${cx(config)} {
+        .model-item-config {
           opacity: 1;
         }
       }

package/src/libs/next/config/define-config.ts

@@ -311,7 +311,7 @@ export function defineConfig(config: CustomNextConfig) {
     ],
 
     // when external packages in dev mode with turbopack, this config will lead to bundle error
-    serverExternalPackages: isProd ? ['@electric-sql/pglite', 'pdfkit'] : ['pdfkit'],
+    serverExternalPackages: ['pdfkit'],
 
     transpilePackages: ['pdfjs-dist', 'mermaid', 'better-auth-harmony'],
     turbopack: {

package/src/server/modules/S3/index.test.ts

@@ -3,6 +3,7 @@ import {
   DeleteObjectCommand,
   DeleteObjectsCommand,
   GetObjectCommand,
+  HeadObjectCommand,
   PutObjectCommand,
   S3Client,
 } from '@aws-sdk/client-s3';
@@ -304,6 +305,63 @@ describe('FileS3', () => {
     });
   });
 
+  describe('getFileMetadata', () => {
+    it('should retrieve file metadata with content length and type', async () => {
+      const s3 = new FileS3();
+      mockS3ClientSend.mockResolvedValue({
+        ContentLength: 1024,
+        ContentType: 'image/png',
+      });
+
+      const result = await s3.getFileMetadata('test-file.png');
+
+      expect(HeadObjectCommand).toHaveBeenCalledWith({
+        Bucket: 'test-bucket',
+        Key: 'test-file.png',
+      });
+      expect(result).toEqual({
+        contentLength: 1024,
+        contentType: 'image/png',
+      });
+    });
+
+    it('should return 0 for content length when not provided', async () => {
+      const s3 = new FileS3();
+      mockS3ClientSend.mockResolvedValue({
+        ContentType: 'application/octet-stream',
+      });
+
+      const result = await s3.getFileMetadata('test-file.bin');
+
+      expect(result).toEqual({
+        contentLength: 0,
+        contentType: 'application/octet-stream',
+      });
+    });
+
+    it('should handle missing content type', async () => {
+      const s3 = new FileS3();
+      mockS3ClientSend.mockResolvedValue({
+        ContentLength: 2048,
+      });
+
+      const result = await s3.getFileMetadata('test-file.bin');
+
+      expect(result).toEqual({
+        contentLength: 2048,
+        contentType: undefined,
+      });
+    });
+
+    it('should handle S3 errors', async () => {
+      const s3 = new FileS3();
+      const error = new Error('File not found');
+      mockS3ClientSend.mockRejectedValue(error);
+
+      await expect(s3.getFileMetadata('non-existent-file.txt')).rejects.toThrow('File not found');
+    });
+  });
+
   describe('createPreSignedUrl', () => {
     it('should create presigned URL for upload with ACL', async () => {
       const s3 = new FileS3();

package/src/server/modules/S3/index.ts

@@ -2,6 +2,7 @@ import {
   DeleteObjectCommand,
   DeleteObjectsCommand,
   GetObjectCommand,
+  HeadObjectCommand,
   PutObjectCommand,
   S3Client,
 } from '@aws-sdk/client-s3';
@@ -111,6 +112,26 @@ export class S3 {
     return response.Body.transformToByteArray();
   }
 
+  /**
+   * Get file metadata from S3 using HeadObject
+   * This is used to verify actual file size from S3 instead of trusting client-provided values
+   */
+  public async getFileMetadata(
+    key: string,
+  ): Promise<{ contentLength: number; contentType?: string }> {
+    const command = new HeadObjectCommand({
+      Bucket: this.bucket,
+      Key: key,
+    });
+
+    const response = await this.client.send(command);
+
+    return {
+      contentLength: response.ContentLength ?? 0,
+      contentType: response.ContentType,
+    };
+  }
+
   public async createPreSignedUrl(key: string): Promise<string> {
     const command = new PutObjectCommand({
       ACL: this.setAcl ? 'public-read' : undefined,