@lobehub/chat 1.98.2 → 1.99.0

package/src/database/models/__tests__/generationTopic.test.ts

@@ -0,0 +1,411 @@
+// @vitest-environment node
+import { eq } from 'drizzle-orm/expressions';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { LobeChatDatabase } from '@/database/type';
+import { FileService } from '@/server/services/file';
+import { ImageGenerationTopic } from '@/types/generation';
+
+import { generationBatches, generationTopics, generations, users } from '../../schemas';
+import { GenerationTopicItem } from '../../schemas/generation';
+import { GenerationTopicModel } from '../generationTopic';
+import { getTestDB } from './_util';
+
+// Mock FileService
+const mockGetFullFileUrl = vi.fn();
+vi.mock('@/server/services/file', () => ({
+  FileService: vi.fn().mockImplementation(() => ({
+    getFullFileUrl: mockGetFullFileUrl,
+  })),
+}));
+
+const serverDB: LobeChatDatabase = await getTestDB();
+
+const userId = 'generation-topic-test-user';
+const otherUserId = 'other-user';
+const generationTopicModel = new GenerationTopicModel(serverDB, userId);
+
+beforeEach(async () => {
+  await serverDB.delete(users);
+  await serverDB.insert(users).values([{ id: userId }, { id: otherUserId }]);
+
+  // Reset mocks before each test
+  vi.clearAllMocks();
+  mockGetFullFileUrl.mockImplementation((url: string) => `https://example.com/${url}`);
+});
+
+afterEach(async () => {
+  await serverDB.delete(users);
+});
+
+describe('GenerationTopicModel', () => {
+  describe('create', () => {
+    it('should create a new generation topic', async () => {
+      const title = 'Test Generation Topic';
+
+      const result = await generationTopicModel.create(title);
+
+      expect(result.id).toBeDefined();
+      expect(result.title).toBe(title);
+      expect(result.userId).toBe(userId);
+      expect(result.coverUrl).toBeNull();
+      expect(result.createdAt).toBeInstanceOf(Date);
+      expect(result.updatedAt).toBeInstanceOf(Date);
+
+      // Verify it's saved in database
+      const topic = await serverDB.query.generationTopics.findFirst({
+        where: eq(generationTopics.id, result.id),
+      });
+      expect(topic).toMatchObject({ title, userId });
+    });
+
+    it('should create a topic with empty title', async () => {
+      const result = await generationTopicModel.create('');
+
+      expect(result.id).toBeDefined();
+      expect(result.title).toBe('');
+      expect(result.userId).toBe(userId);
+    });
+  });
+
+  describe('queryAll', () => {
+    it('should return all topics for the user ordered by updatedAt desc', async () => {
+      // Create test data with different timestamps
+      const now = new Date();
+      const earlier = new Date(now.getTime() - 60000); // 1 minute earlier
+      const earliest = new Date(now.getTime() - 120000); // 2 minutes earlier
+
+      await serverDB.insert(generationTopics).values([
+        {
+          id: 'topic1',
+          userId,
+          title: 'Topic 1',
+          updatedAt: earliest,
+        },
+        {
+          id: 'topic2',
+          userId,
+          title: 'Topic 2',
+          updatedAt: now,
+        },
+        {
+          id: 'topic3',
+          userId,
+          title: 'Topic 3',
+          updatedAt: earlier,
+        },
+        {
+          id: 'topic4',
+          userId: otherUserId,
+          title: 'Other User Topic',
+          updatedAt: now,
+        },
+      ]);
+
+      const result = await generationTopicModel.queryAll();
+
+      // Should return only topics for current user, ordered by updatedAt desc
+      expect(result).toHaveLength(3);
+      expect(result[0].id).toBe('topic2'); // most recent
+      expect(result[1].id).toBe('topic3'); // middle
+      expect(result[2].id).toBe('topic1'); // oldest
+    });
+
+    it('should process cover URLs through FileService', async () => {
+      await serverDB.insert(generationTopics).values([
+        {
+          id: 'topic1',
+          userId,
+          title: 'Topic with cover',
+          coverUrl: 'cover-image-key',
+        },
+        {
+          id: 'topic2',
+          userId,
+          title: 'Topic without cover',
+          coverUrl: null,
+        },
+      ]);
+
+      const result = await generationTopicModel.queryAll();
+
+      expect(result).toHaveLength(2);
+      expect(result[0].coverUrl).toBe('https://example.com/cover-image-key');
+      expect(result[1].coverUrl).toBeNull();
+
+      // Verify FileService was called for the topic with coverUrl
+      expect(mockGetFullFileUrl).toHaveBeenCalledWith('cover-image-key');
+      expect(mockGetFullFileUrl).toHaveBeenCalledTimes(1);
+    });
+
+    it('should return empty array if no topics exist', async () => {
+      const result = await generationTopicModel.queryAll();
+      expect(result).toHaveLength(0);
+    });
+  });
+
+  describe('update', () => {
+    it('should update a generation topic', async () => {
+      // Create a test topic
+      const { id } = await generationTopicModel.create('Original Title');
+
+      const updateData: Partial<ImageGenerationTopic> = {
+        title: 'Updated Title',
+        coverUrl: 'new-cover-key',
+      };
+
+      const result = await generationTopicModel.update(id, updateData);
+
+      expect(result).toBeDefined();
+      expect(result!.id).toBe(id);
+      expect(result!.title).toBe('Updated Title');
+      expect(result!.coverUrl).toBe('new-cover-key');
+      expect(result!.updatedAt).toBeInstanceOf(Date);
+
+      // Verify in database
+      const updatedTopic = await serverDB.query.generationTopics.findFirst({
+        where: eq(generationTopics.id, id),
+      });
+      expect(updatedTopic).toMatchObject(updateData);
+    });
+
+    it('should not update topics of other users', async () => {
+      // Create a topic for another user
+      const [otherUserTopic] = await serverDB
+        .insert(generationTopics)
+        .values({ id: 'other-topic', userId: otherUserId, title: 'Other User Topic' })
+        .returning();
+
+      const updateData: Partial<ImageGenerationTopic> = {
+        title: 'Hacked Title',
+      };
+
+      // Attempt to update should not affect other user's topic
+      const result = await generationTopicModel.update(otherUserTopic.id, updateData);
+
+      // Should return undefined or empty result because of user permission check
+      expect(result).toBeUndefined();
+
+      // Verify the topic remains unchanged in the database
+      const unchangedTopic = await serverDB.query.generationTopics.findFirst({
+        where: eq(generationTopics.id, otherUserTopic.id),
+      });
+      expect(unchangedTopic?.title).toBe('Other User Topic');
+    });
+
+    it('should update only specified fields', async () => {
+      const { id } = await generationTopicModel.create('Original Title');
+
+      // Update only title
+      const result = await generationTopicModel.update(id, { title: 'Only Title Updated' });
+
+      expect(result).toBeDefined();
+      expect(result!.title).toBe('Only Title Updated');
+      expect(result!.coverUrl).toBeNull(); // Should remain unchanged
+    });
+  });
+
+  describe('delete', () => {
+    it('should delete a generation topic', async () => {
+      const { id } = await generationTopicModel.create('Topic to Delete');
+
+      const result = await generationTopicModel.delete(id);
+
+      expect(result).toBeDefined();
+      const deleteResult = result!;
+      expect(deleteResult.deletedTopic.id).toBe(id);
+      expect(deleteResult.deletedTopic.title).toBe('Topic to Delete');
+      expect(deleteResult.filesToDelete).toEqual([]);
+
+      // Verify it's deleted from database
+      const deletedTopic = await serverDB.query.generationTopics.findFirst({
+        where: eq(generationTopics.id, id),
+      });
+      expect(deletedTopic).toBeUndefined();
+    });
+
+    it('should not delete topics of other users', async () => {
+      // Create a topic for another user
+      const [otherUserTopic] = await serverDB
+        .insert(generationTopics)
+        .values({ id: 'other-topic', userId: otherUserId, title: 'Other User Topic' })
+        .returning();
+
+      // Attempt to delete should not affect other user's topic
+      const result = await generationTopicModel.delete(otherUserTopic.id);
+
+      // Should return undefined because of user permission check
+      expect(result).toBeUndefined();
+
+      // The topic should still exist in the database
+      const stillExistsTopic = await serverDB.query.generationTopics.findFirst({
+        where: eq(generationTopics.id, otherUserTopic.id),
+      });
+      expect(stillExistsTopic).toBeDefined();
+      expect(stillExistsTopic?.title).toBe('Other User Topic');
+    });
+
+    it('should return deleted topic data', async () => {
+      const { id } = await generationTopicModel.create('Topic with Data');
+
+      // Add some data to the topic
+      await generationTopicModel.update(id, {
+        title: 'Updated Topic',
+        coverUrl: 'cover-key',
+      });
+
+      const result = await generationTopicModel.delete(id);
+
+      expect(result).toBeDefined();
+      const deleteResult = result!;
+      expect(deleteResult.deletedTopic).toMatchObject({
+        id,
+        title: 'Updated Topic',
+        coverUrl: 'cover-key',
+        userId,
+      });
+      expect(deleteResult.filesToDelete).toContain('cover-key');
+    });
+
+    it('should return undefined when trying to delete non-existent topic', async () => {
+      const nonExistentId = 'non-existent-topic-id';
+
+      const result = await generationTopicModel.delete(nonExistentId);
+
+      // Should return undefined because topic doesn't exist
+      expect(result).toBeUndefined();
+    });
+
+    it('should return undefined when trying to delete topic with invalid format id', async () => {
+      const invalidId = 'invalid-format-id';
+
+      const result = await generationTopicModel.delete(invalidId);
+
+      // Should return undefined because topic doesn't exist with this invalid ID
+      expect(result).toBeUndefined();
+    });
+
+    it('should collect file URLs from batches and generations when deleting topic with data', async () => {
+      // Create a topic with cover image
+      const { id: topicId } = await generationTopicModel.create(
+        'Topic with batches and generations',
+      );
+      await generationTopicModel.update(topicId, { coverUrl: 'topic-cover.jpg' });
+
+      // Create a generation batch associated with this topic
+      const [batch] = await serverDB
+        .insert(generationBatches)
+        .values({
+          userId,
+          generationTopicId: topicId,
+          provider: 'test-provider',
+          model: 'test-model',
+          prompt: 'Test generation prompt',
+          width: 1024,
+          height: 1024,
+        })
+        .returning();
+
+      // Create generations with asset data containing thumbnail URLs
+      await serverDB.insert(generations).values([
+        {
+          userId,
+          generationBatchId: batch.id,
+          asyncTaskId: null,
+          fileId: null,
+          seed: 12345,
+          asset: {
+            type: 'image',
+            thumbnailUrl: 'thumbnail1.jpg',
+            originalUrl: 'original1.jpg',
+            width: 1024,
+            height: 1024,
+          },
+        },
+        {
+          userId,
+          generationBatchId: batch.id,
+          asyncTaskId: null,
+          fileId: null,
+          seed: 12346,
+          asset: {
+            type: 'image',
+            thumbnailUrl: 'thumbnail2.jpg',
+            originalUrl: 'original2.jpg',
+            width: 1024,
+            height: 1024,
+          },
+        },
+      ]);
+
+      // Now delete the topic - this should collect all file URLs from cover + generations
+      const result = await generationTopicModel.delete(topicId);
+
+      expect(result).toBeDefined();
+      const deleteResult = result!;
+      expect(deleteResult.deletedTopic.id).toBe(topicId);
+
+      // Should collect cover URL and thumbnail URLs from generations (lines 111-117)
+      expect(deleteResult.filesToDelete).toContain('topic-cover.jpg');
+      expect(deleteResult.filesToDelete).toContain('thumbnail1.jpg');
+      expect(deleteResult.filesToDelete).toContain('thumbnail2.jpg');
+      expect(deleteResult.filesToDelete).toHaveLength(3);
+
+      // Verify topic is actually deleted from database
+      const deletedTopic = await serverDB.query.generationTopics.findFirst({
+        where: eq(generationTopics.id, topicId),
+      });
+      expect(deletedTopic).toBeUndefined();
+    });
+  });
+
+  describe('user isolation', () => {
+    it('should only operate on topics belonging to the user', async () => {
+      // Create topics for different users
+      await serverDB.insert(generationTopics).values([
+        { id: 'user1-topic1', userId, title: 'User 1 Topic 1' },
+        { id: 'user1-topic2', userId, title: 'User 1 Topic 2' },
+        { id: 'user2-topic1', userId: otherUserId, title: 'User 2 Topic 1' },
+      ]);
+
+      const result = await generationTopicModel.queryAll();
+
+      // Should only return topics for the current user
+      expect(result).toHaveLength(2);
+      expect(result.every((topic) => topic.userId === userId)).toBe(true);
+      expect(result.some((topic) => topic.title === 'User 2 Topic 1')).toBe(false);
+    });
+  });
+
+  describe('edge cases', () => {
+    it('should handle topics with null titles', async () => {
+      await serverDB.insert(generationTopics).values({
+        id: 'null-title-topic',
+        userId,
+        title: null,
+      });
+
+      const result = await generationTopicModel.queryAll();
+
+      expect(result).toHaveLength(1);
+      expect(result[0].title).toBeNull();
+    });
+
+    it('should handle topics with null coverUrl', async () => {
+      const { id } = await generationTopicModel.create('Topic');
+
+      const result = await generationTopicModel.update(id, { coverUrl: null });
+
+      expect(result).toBeDefined();
+      expect(result!.coverUrl).toBeNull();
+    });
+
+    it('should return undefined when updating non-existent topic', async () => {
+      const nonExistentId = 'non-existent-topic';
+
+      const updateResult = await generationTopicModel.update(nonExistentId, { title: 'New Title' });
+
+      expect(updateResult).toBeUndefined();
+    });
+  });
+});

package/src/database/models/aiModel.ts

@@ -67,6 +67,7 @@ export class AiModelModel {
         displayName: aiModels.displayName,
         enabled: aiModels.enabled,
         id: aiModels.id,
+        parameters: aiModels.parameters,
         pricing: aiModels.pricing,
         releasedAt: aiModels.releasedAt,
         source: aiModels.source,
@@ -93,6 +94,7 @@ export class AiModelModel {
         displayName: aiModels.displayName,
         enabled: aiModels.enabled,
         id: aiModels.id,
+        parameters: aiModels.parameters,
         providerId: aiModels.providerId,
         sort: aiModels.sort,
         source: aiModels.source,

package/src/database/models/asyncTask.ts

@@ -82,7 +82,7 @@ export class AsyncTaskModel {
         .set({
           error: new AsyncTaskError(
             AsyncTaskErrorType.Timeout,
-            'chunking task is timeout, please try again',
+            'task is timeout, please try again',
           ),
           status: AsyncTaskStatus.Error,
         })

package/src/database/models/file.ts

@@ -2,7 +2,7 @@ import { count, sum } from 'drizzle-orm';
 import { and, asc, desc, eq, ilike, inArray, like, notExists, or } from 'drizzle-orm/expressions';
 import type { PgTransaction } from 'drizzle-orm/pg-core';
 
-import { LobeChatDatabase } from '@/database/type';
+import { LobeChatDatabase, Transaction } from '@/database/type';
 import { FilesTabs, QueryFileListParams, SortType } from '@/types/files';
 
 import {
@@ -29,10 +29,11 @@ export class FileModel {
   create = async (
     params: Omit<NewFile, 'id' | 'userId'> & { knowledgeBaseId?: string },
     insertToGlobalFiles?: boolean,
+    trx?: Transaction,
   ) => {
-    const result = await this.db.transaction(async (trx) => {
+    const executeInTransaction = async (tx: Transaction) => {
       if (insertToGlobalFiles) {
-        await trx.insert(globalFiles).values({
+        await tx.insert(globalFiles).values({
           creator: this.userId,
           fileType: params.fileType,
           hashId: params.fileHash!,
@@ -42,7 +43,7 @@ export class FileModel {
         });
       }
 
-      const result = await trx
+      const result = await tx
         .insert(files)
         .values({ ...params, userId: this.userId })
         .returning();
@@ -50,7 +51,7 @@ export class FileModel {
       const item = result[0];
 
       if (params.knowledgeBaseId) {
-        await trx.insert(knowledgeBaseFiles).values({
+        await tx.insert(knowledgeBaseFiles).values({
           fileId: item.id,
           knowledgeBaseId: params.knowledgeBaseId,
           userId: this.userId,
@@ -58,8 +59,11 @@ export class FileModel {
       }
 
       return item;
-    });
+    };
 
+    const result = await (trx
+      ? executeInTransaction(trx)
+      : this.db.transaction(executeInTransaction));
     return { id: result.id };
   };
 
@@ -82,20 +86,21 @@ export class FileModel {
     };
   };
 
-  delete = async (id: string, removeGlobalFile: boolean = true) => {
-    const file = await this.findById(id);
-    if (!file) return;
+  delete = async (id: string, removeGlobalFile: boolean = true, trx?: Transaction) => {
+    const executeInTransaction = async (tx: Transaction) => {
+      // pglite 环境下不能再 transaction 中使用非事务操作，会阻塞住
+      const file = await this.findById(id, tx);
+      if (!file) return;
 
-    const fileHash = file.fileHash!;
+      const fileHash = file.fileHash!;
 
-    return await this.db.transaction(async (trx) => {
-      // 1. 删除相关的 chunks
-      await this.deleteFileChunks(trx as any, [id]);
+      // 2. Delete related chunks
+      await this.deleteFileChunks(tx as any, [id]);
 
-      // 2. 删除文件记录
-      await trx.delete(files).where(and(eq(files.id, id), eq(files.userId, this.userId)));
+      // 3. Delete file record
+      await tx.delete(files).where(and(eq(files.id, id), eq(files.userId, this.userId)));
 
-      const result = await trx
+      const result = await tx
         .select({ count: count() })
         .from(files)
         .where(and(eq(files.fileHash, fileHash)));
@@ -105,11 +110,13 @@ export class FileModel {
       // delete the file from global file if it is not used by other files
       // if `DISABLE_REMOVE_GLOBAL_FILE` is true, we will not remove the global file
       if (fileCount === 0 && removeGlobalFile) {
-        await trx.delete(globalFiles).where(eq(globalFiles.hashId, fileHash));
+        await tx.delete(globalFiles).where(eq(globalFiles.hashId, fileHash));
 
         return file;
       }
-    });
+    };
+
+    return await (trx ? executeInTransaction(trx) : this.db.transaction(executeInTransaction));
   };
 
   deleteGlobalFile = async (hashId: string) => {
@@ -261,8 +268,9 @@ export class FileModel {
     });
   };
 
-  findById = async (id: string) => {
-    return this.db.query.files.findFirst({
+  findById = async (id: string, trx?: Transaction) => {
+    const database = trx || this.db;
+    return database.query.files.findFirst({
       where: and(eq(files.id, id), eq(files.userId, this.userId)),
     });
   };