@lobehub/lobehub 2.0.0-next.137 → 2.0.0-next.139

package/packages/conversation-flow/src/transformation/__tests__/ContextTreeBuilder.test.ts

@@ -311,6 +311,70 @@ describe('ContextTreeBuilder', () => {
       expect(result).toHaveLength(0);
     });
 
+    it('should set activeBranchIndex to children.length for optimistic update', () => {
+      const messageMap = new Map<string, Message>([
+        [
+          'msg-1',
+          {
+            content: 'Hello',
+            createdAt: 0,
+            id: 'msg-1',
+            meta: {},
+            // activeBranchIndex = 2 means optimistic update (pointing to not-yet-created branch)
+            metadata: { activeBranchIndex: 2 },
+            role: 'user',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-2',
+          {
+            content: 'Response 1',
+            createdAt: 0,
+            id: 'msg-2',
+            meta: {},
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-3',
+          {
+            content: 'Response 2',
+            createdAt: 0,
+            id: 'msg-3',
+            meta: {},
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+      ]);
+
+      const builder = createBuilder(messageMap);
+      const idNodes: IdNode[] = [
+        {
+          children: [
+            { children: [], id: 'msg-2' },
+            { children: [], id: 'msg-3' },
+          ],
+          id: 'msg-1',
+        },
+      ];
+
+      const result = builder.transformAll(idNodes);
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toEqual({ id: 'msg-1', type: 'message' });
+      // When activeBranchIndex === children.length (optimistic update),
+      // BranchResolver returns undefined, and ContextTreeBuilder uses children.length as index
+      expect(result[1]).toMatchObject({
+        activeBranchIndex: 2, // children.length = 2
+        branches: [[{ id: 'msg-2', type: 'message' }], [{ id: 'msg-3', type: 'message' }]],
+        parentMessageId: 'msg-1',
+        type: 'branch',
+      });
+    });
+
     it('should continue with active column children in compare mode', () => {
       const messageMap = new Map<string, Message>([
         [

package/packages/conversation-flow/src/transformation/__tests__/FlatListBuilder.test.ts

@@ -507,5 +507,52 @@ describe('FlatListBuilder', () => {
       expect(result[1].role).toBe('assistantGroup');
       expect(result[2].id).toBe('msg-3');
     });
+
+    it('should handle optimistic update for user message with branches', () => {
+      // Scenario: User has sent a new message, activeBranchIndex points to a branch
+      // that is being created but doesn't exist yet (optimistic update)
+      const messages: Message[] = [
+        {
+          content: 'User',
+          createdAt: 0,
+          id: 'msg-1',
+          meta: {},
+          // activeBranchIndex = 2 means pointing to a not-yet-created branch (optimistic update)
+          // when there are only 2 existing children (msg-2, msg-3)
+          metadata: { activeBranchIndex: 2 },
+          role: 'user',
+          updatedAt: 0,
+        },
+        {
+          content: 'Branch 1',
+          createdAt: 0,
+          id: 'msg-2',
+          meta: {},
+          parentId: 'msg-1',
+          role: 'assistant',
+          updatedAt: 0,
+        },
+        {
+          content: 'Branch 2',
+          createdAt: 0,
+          id: 'msg-3',
+          meta: {},
+          parentId: 'msg-1',
+          role: 'assistant',
+          updatedAt: 0,
+        },
+      ];
+
+      const builder = createBuilder(messages);
+      const result = builder.flatten(messages);
+
+      // When activeBranchIndex === children.length (optimistic update),
+      // BranchResolver returns undefined, and FlatListBuilder just adds the user message
+      // without branch info and doesn't continue to any branch
+      expect(result).toHaveLength(1);
+      expect(result[0].id).toBe('msg-1');
+      // User message should not have branch info since we're in optimistic update mode
+      expect((result[0] as any).siblingCount).toBeUndefined();
+    });
   });
 });

package/packages/database/src/models/__tests__/agent.test.ts

@@ -20,9 +20,12 @@ import { getTestDB } from './_util';
 const serverDB: LobeChatDatabase = await getTestDB();
 
 const userId = 'agent-model-test-user-id';
+const userId2 = 'agent-model-test-user-id-2';
 const agentModel = new AgentModel(serverDB, userId);
+const agentModel2 = new AgentModel(serverDB, userId2);
 
 const knowledgeBase = { id: 'kb1', userId, name: 'knowledgeBase' };
+const knowledgeBase2 = { id: 'kb2', userId: userId2, name: 'knowledgeBase2' };
 const fileList = [
   {
     id: '1',
@@ -42,11 +45,22 @@ const fileList = [
   },
 ];
 
+const fileList2 = [
+  {
+    id: '3',
+    name: 'other.pdf',
+    url: 'https://a.com/other.pdf',
+    size: 1000,
+    fileType: 'application/pdf',
+    userId: userId2,
+  },
+];
+
 beforeEach(async () => {
   await serverDB.delete(users);
-  await serverDB.insert(users).values([{ id: userId }]);
-  await serverDB.insert(knowledgeBases).values(knowledgeBase);
-  await serverDB.insert(files).values(fileList);
+  await serverDB.insert(users).values([{ id: userId }, { id: userId2 }]);
+  await serverDB.insert(knowledgeBases).values([knowledgeBase, knowledgeBase2]);
+  await serverDB.insert(files).values([...fileList, ...fileList2]);
 });
 
 afterEach(async () => {
@@ -226,6 +240,27 @@ describe('AgentModel', () => {
 
       expect(result).toBeUndefined();
     });
+
+    it('should not delete another user agent knowledge base association', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+      await serverDB
+        .insert(agentsKnowledgeBases)
+        .values({ agentId: agent.id, knowledgeBaseId: knowledgeBase.id, userId });
+
+      // Try to delete with another user's model
+      await agentModel2.deleteAgentKnowledgeBase(agent.id, knowledgeBase.id);
+
+      const result = await serverDB.query.agentsKnowledgeBases.findFirst({
+        where: eq(agentsKnowledgeBases.agentId, agent.id),
+      });
+
+      // Should still exist
+      expect(result).toBeDefined();
+    });
   });
 
   describe('toggleKnowledgeBase', () => {
@@ -248,6 +283,28 @@ describe('AgentModel', () => {
 
       expect(result?.enabled).toBe(false);
     });
+
+    it('should not toggle another user agent knowledge base association', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+
+      await serverDB
+        .insert(agentsKnowledgeBases)
+        .values({ agentId: agent.id, knowledgeBaseId: knowledgeBase.id, userId, enabled: true });
+
+      // Try to toggle with another user's model
+      await agentModel2.toggleKnowledgeBase(agent.id, knowledgeBase.id, false);
+
+      const result = await serverDB.query.agentsKnowledgeBases.findFirst({
+        where: eq(agentsKnowledgeBases.agentId, agent.id),
+      });
+
+      // Should still be enabled
+      expect(result?.enabled).toBe(true);
+    });
   });
 
   describe('createAgentFiles', () => {
@@ -363,6 +420,26 @@ describe('AgentModel', () => {
 
       expect(result).toBeUndefined();
     });
+
+    it('should not delete another user agent file association', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+
+      await serverDB.insert(agentsFiles).values({ agentId: agent.id, fileId: '1', userId });
+
+      // Try to delete with another user's model
+      await agentModel2.deleteAgentFile(agent.id, '1');
+
+      const result = await serverDB.query.agentsFiles.findFirst({
+        where: eq(agentsFiles.agentId, agent.id),
+      });
+
+      // Should still exist
+      expect(result).toBeDefined();
+    });
   });
 
   describe('toggleFile', () => {
@@ -385,5 +462,27 @@ describe('AgentModel', () => {
 
       expect(result?.enabled).toBe(false);
     });
+
+    it('should not toggle another user agent file association', async () => {
+      const agent = await serverDB
+        .insert(agents)
+        .values({ userId })
+        .returning()
+        .then((res) => res[0]);
+
+      await serverDB
+        .insert(agentsFiles)
+        .values({ agentId: agent.id, fileId: '1', userId, enabled: true });
+
+      // Try to toggle with another user's model
+      await agentModel2.toggleFile(agent.id, '1', false);
+
+      const result = await serverDB.query.agentsFiles.findFirst({
+        where: eq(agentsFiles.agentId, agent.id),
+      });
+
+      // Should still be enabled
+      expect(result?.enabled).toBe(true);
+    });
   });
 });

package/packages/database/src/models/__tests__/document.test.ts

@@ -54,6 +54,169 @@ const createTestDocument = async (model: DocumentModel, fModel: FileModel, conte
 };
 
 describe('DocumentModel', () => {
+  describe('create', () => {
+    it('should create a new document', async () => {
+      const { id: fileId } = await fileModel.create({
+        fileType: 'text/plain',
+        name: 'test.txt',
+        size: 100,
+        url: 'https://example.com/test.txt',
+      });
+
+      const file = await fileModel.findById(fileId);
+      if (!file) throw new Error('File not found');
+
+      const result = await documentModel.create({
+        content: 'Test content',
+        fileId: file.id,
+        fileType: 'text/plain',
+        source: file.url,
+        sourceType: 'file',
+        totalCharCount: 12,
+        totalLineCount: 1,
+      });
+
+      expect(result).toBeDefined();
+      expect(result.content).toBe('Test content');
+      expect(result.fileId).toBe(file.id);
+    });
+  });
+
+  describe('delete', () => {
+    it('should delete a document', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content');
+
+      await documentModel.delete(documentId);
+
+      const deleted = await documentModel.findById(documentId);
+      expect(deleted).toBeUndefined();
+    });
+
+    it('should not delete document belonging to another user', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content');
+
+      // Try to delete with another user's model
+      await documentModel2.delete(documentId);
+
+      // Document should still exist
+      const stillExists = await documentModel.findById(documentId);
+      expect(stillExists).toBeDefined();
+    });
+  });
+
+  describe('deleteAll', () => {
+    it('should delete all documents for the user', async () => {
+      await createTestDocument(documentModel, fileModel, 'First document');
+      await createTestDocument(documentModel, fileModel, 'Second document');
+      await createTestDocument(documentModel2, fileModel2, 'Other user document');
+
+      await documentModel.deleteAll();
+
+      const userDocs = await documentModel.query();
+      const otherUserDocs = await documentModel2.query();
+
+      expect(userDocs).toHaveLength(0);
+      expect(otherUserDocs).toHaveLength(1);
+    });
+  });
+
+  describe('query', () => {
+    it('should return all documents for the user', async () => {
+      await createTestDocument(documentModel, fileModel, 'First document');
+      await createTestDocument(documentModel, fileModel, 'Second document');
+
+      const result = await documentModel.query();
+
+      expect(result).toHaveLength(2);
+    });
+
+    it('should only return documents for the current user', async () => {
+      await createTestDocument(documentModel, fileModel, 'User 1 document');
+      await createTestDocument(documentModel2, fileModel2, 'User 2 document');
+
+      const result = await documentModel.query();
+
+      expect(result).toHaveLength(1);
+      expect(result[0].content).toBe('User 1 document');
+    });
+
+    it('should return documents ordered by updatedAt desc', async () => {
+      const { documentId: doc1Id } = await createTestDocument(
+        documentModel,
+        fileModel,
+        'First document',
+      );
+      const { documentId: doc2Id } = await createTestDocument(
+        documentModel,
+        fileModel,
+        'Second document',
+      );
+
+      // Wait a bit to ensure timestamp difference
+      await new Promise((resolve) => setTimeout(resolve, 50));
+
+      // Update first document to make it more recent
+      await documentModel.update(doc1Id, { content: 'Updated first document' });
+
+      const result = await documentModel.query();
+
+      expect(result[0].id).toBe(doc1Id);
+      expect(result[1].id).toBe(doc2Id);
+    });
+  });
+
+  describe('findById', () => {
+    it('should find document by id', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content');
+
+      const found = await documentModel.findById(documentId);
+
+      expect(found).toBeDefined();
+      expect(found?.id).toBe(documentId);
+      expect(found?.content).toBe('Test content');
+    });
+
+    it('should return undefined for non-existent document', async () => {
+      const found = await documentModel.findById('non-existent-id');
+
+      expect(found).toBeUndefined();
+    });
+
+    it('should not find document belonging to another user', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content');
+
+      const found = await documentModel2.findById(documentId);
+
+      expect(found).toBeUndefined();
+    });
+  });
+
+  describe('update', () => {
+    it('should update a document', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Original content');
+
+      await documentModel.update(documentId, {
+        content: 'Updated content',
+        totalCharCount: 15,
+      });
+
+      const updated = await documentModel.findById(documentId);
+
+      expect(updated?.content).toBe('Updated content');
+      expect(updated?.totalCharCount).toBe(15);
+    });
+
+    it('should not update document belonging to another user', async () => {
+      const { documentId } = await createTestDocument(documentModel, fileModel, 'Original content');
+
+      await documentModel2.update(documentId, { content: 'Hacked content' });
+
+      const unchanged = await documentModel.findById(documentId);
+
+      expect(unchanged?.content).toBe('Original content');
+    });
+  });
+
   describe('findByFileId', () => {
     it('should find document by fileId', async () => {
       const { documentId, file } = await createTestDocument(