chore(server): ignore non-exists doc embedding (#11153)

.github/actions/server-test-env/action.yml

@@ -19,5 +19,5 @@ runs:
         NODE_ENV: test
       run: |
         yarn affine @affine/server prisma generate
-        yarn affine @affine/server prisma migrate reset -f
+        yarn affine @affine/server prisma migrate deploy
         yarn affine @affine/server data-migration run

packages/backend/server/src/__tests__/copilot.spec.ts

@@ -1381,19 +1381,24 @@ test('should be able to manage context', async t => {
     }
 
     // doc record
-    const docId = randomUUID();
+
-    await t.context.db.snapshot.create({
+    const addDoc = async () => {
-      data: {
+      const docId = randomUUID();
-        workspaceId: session.workspaceId,
+      await t.context.db.snapshot.create({
-        id: docId,
+        data: {
-        blob: Buffer.from([1, 1]),
+          workspaceId: session.workspaceId,
-        state: Buffer.from([1, 1]),
+          id: docId,
-        updatedAt: new Date(),
+          blob: Buffer.from([1, 1]),
-        createdAt: new Date(),
+          state: Buffer.from([1, 1]),
-      },
+          updatedAt: new Date(),
-    });
+          createdAt: new Date(),
+        },
+      });
+      return docId;
+    };
 
     {
+      const docId = await addDoc();
       await session.addDocRecord(docId);
       const docs = session.docs.map(d => d.id);
       t.deepEqual(docs, [docId], 'should list doc id');
@@ -1405,50 +1410,65 @@ test('should be able to manage context', async t => {
     // tag record
     {
       const tagId = randomUUID();
-      await session.addCategoryRecord(ContextCategories.Tag, tagId, [docId]);
+
-      const tags = session.tags.map(t => t.id);
+      const docId1 = await addDoc();
-      t.deepEqual(tags, [tagId], 'should list tag id');
+      const docId2 = await addDoc();
+
+      {
+        await session.addCategoryRecord(ContextCategories.Tag, tagId, [docId1]);
+        const tags = session.tags.map(t => t.id);
+        t.deepEqual(tags, [tagId], 'should list tag id');
+
+        const docs = session.tags.flatMap(l => l.docs.map(d => d.id));
+        t.deepEqual(docs, [docId1], 'should list doc ids');
+      }
+
+      {
+        await session.addCategoryRecord(ContextCategories.Tag, tagId, [docId2]);
+
+        const docs = session.tags.flatMap(l => l.docs.map(d => d.id));
+        t.deepEqual(docs, [docId1, docId2], 'should list doc ids');
+      }
 
       await session.removeCategoryRecord(ContextCategories.Tag, tagId);
       t.deepEqual(session.tags, [], 'should remove tag id');
-
-      await t.throwsAsync(
-        session.addCategoryRecord(ContextCategories.Tag, tagId, [
-          'not-exists-doc',
-        ]),
-        {
-          instanceOf: Error,
-        },
-        'should throw error if doc id not exists'
-      );
     }
 
     // collection record
     {
       const collectionId = randomUUID();
-      await session.addCategoryRecord(
+
-        ContextCategories.Collection,
+      const docId1 = await addDoc();
-        collectionId,
+      const docId2 = await addDoc();
-        [docId]
+      {
-      );
+        await session.addCategoryRecord(
-      const collection = session.collections.map(l => l.id);
+          ContextCategories.Collection,
-      t.deepEqual(collection, [collectionId], 'should list collection id');
+          collectionId,
+          [docId1]
+        );
+        const collection = session.collections.map(l => l.id);
+        t.deepEqual(collection, [collectionId], 'should list collection id');
+
+        const docs = session.collections.flatMap(l => l.docs.map(d => d.id));
+        t.deepEqual(docs, [docId1], 'should list doc ids');
+      }
+
+      {
+        await session.addCategoryRecord(
+          ContextCategories.Collection,
+          collectionId,
+          [docId2]
+        );
+
+        const docs = session.collections.flatMap(l => l.docs.map(d => d.id));
+        t.deepEqual(docs, [docId1, docId2], 'should list doc ids');
+      }
 
       await session.removeCategoryRecord(
         ContextCategories.Collection,
         collectionId
       );
       t.deepEqual(session.collections, [], 'should remove collection id');
-
-      await t.throwsAsync(
-        session.addCategoryRecord(ContextCategories.Collection, collectionId, [
-          'not-exists-doc',
-        ]),
-        {
-          instanceOf: Error,
-        },
-        'should throw error if doc id not exists'
-      );
     }
   }
 });

packages/backend/server/src/plugins/copilot/context/session.ts

@@ -1,6 +1,5 @@
 import { nanoid } from 'nanoid';
 
-import { CopilotDocsNotFound } from '../../../base';
 import {
   ContextCategories,
   ContextCategory,
@@ -63,22 +62,35 @@ export class ContextSession implements AsyncDisposable {
   }
 
   async addCategoryRecord(type: ContextCategories, id: string, docs: string[]) {
-    const existDocs = await this.models.doc.existsAll(this.workspaceId, docs);
-    if (!existDocs) {
-      throw new CopilotDocsNotFound();
-    }
-
     const category = this.config.categories.find(
       c => c.type === type && c.id === id
     );
     if (category) {
+      const missingDocs = docs.filter(
+        docId => !category.docs.some(d => d.id === docId)
+      );
+      if (missingDocs.length) {
+        category.docs.push(
+          ...missingDocs.map(id => ({
+            id,
+            createdAt: Date.now(),
+            status: ContextEmbedStatus.processing,
+          }))
+        );
+        await this.save();
+      }
+
       return category;
     }
     const createdAt = Date.now();
     const record = {
       id,
       type,
-      docs: docs.map(id => ({ id, createdAt, status: null })),
+      docs: docs.map(id => ({
+        id,
+        createdAt,
+        status: ContextEmbedStatus.processing,
+      })),
       createdAt,
     };
     this.config.categories.push(record);