diff --git a/packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.md b/packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.md index 9c1693aefe..f5b84bfb9a 100644 --- a/packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.md +++ b/packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.md @@ -52,6 +52,7 @@ Generated by [AVA](https://avajs.dev). [ { id: 'docId1', + status: 'processing', }, ] diff --git a/packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.snap b/packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.snap index 3bd449d93d..f69a3815f8 100644 Binary files a/packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.snap and b/packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.snap differ diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.md b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.md index 9b90fb7b0f..69bfd67a9a 100644 --- a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.md +++ b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.md @@ -53,3 +53,195 @@ Generated by [AVA](https://avajs.dev). > should return true when embedding table is available true + +## should merge doc status correctly + +> basic doc status merge + + [ + { + id: 'doc1', + status: 'processing', + }, + { + id: 'doc2', + status: 'processing', + }, + { + id: 'doc3', + status: 'failed', + }, + { + id: 'doc4', + status: 'processing', + }, + ] + +> mixed doc status merge + + [ + { + id: 'doc5', + status: 'finished', + }, + { + id: 'doc5', + status: 'finished', + }, + { + id: 'doc6', + status: 'processing', + }, + { + id: 'doc6', + status: 'failed', + }, + { + id: 'doc7', + status: 'processing', + }, + ] + +> edge cases results + + [ + { + case: 0, + length: 1, + statuses: [ + 'processing', + ], + }, + { + case: 1, + length: 1, + statuses: [ + 'processing', + ], + }, + { + case: 2, + length: 100, + statuses: [ + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + 'processing', + ], + }, + ] + +## should handle concurrent mergeDocStatus calls + +> concurrent calls results + + [ + { + call: 1, + status: 'finished', + }, + { + call: 2, + status: 'finished', + }, + { + call: 3, + status: 'processing', + }, + ] diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.snap b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.snap index 91e2cc47e2..661d0e2d34 100644 Binary files a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.snap and b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.snap differ diff --git a/packages/backend/server/src/__tests__/models/copilot-context.spec.ts b/packages/backend/server/src/__tests__/models/copilot-context.spec.ts index 397fd7dfa8..e189f9a9ac 100644 --- a/packages/backend/server/src/__tests__/models/copilot-context.spec.ts +++ b/packages/backend/server/src/__tests__/models/copilot-context.spec.ts @@ -2,8 +2,10 @@ import { randomUUID } from 'node:crypto'; import { AiSession, PrismaClient, User, Workspace } from '@prisma/client'; import ava, { TestFn } from 'ava'; +import Sinon from 'sinon'; import { Config } from '../../base'; +import { ContextEmbedStatus } from '../../models/common/copilot'; import { CopilotContextModel } from '../../models/copilot-context'; import { CopilotSessionModel } from '../../models/copilot-session'; import { CopilotWorkspaceConfigModel } from '../../models/copilot-workspace'; @@ -236,3 +238,173 @@ test('should check embedding table', async t => { // t.false(ret, 'should return false when embedding table is not available'); // } }); + +test('should merge doc status correctly', async t => { + const createDoc = (id: string, status?: string) => ({ + id, + createdAt: Date.now(), + ...(status && { status: status as any }), + }); + + const createDocWithEmbedding = async (docId: string) => { + await t.context.db.snapshot.create({ + data: { + workspaceId: workspace.id, + id: docId, + blob: Buffer.from([1, 1]), + state: Buffer.from([1, 1]), + updatedAt: new Date(), + createdAt: new Date(), + }, + }); + + await t.context.copilotContext.insertWorkspaceEmbedding( + workspace.id, + docId, + [ + { + index: 0, + content: 'content', + embedding: Array.from({ length: 1024 }, () => 1), + }, + ] + ); + }; + + const emptyResult = await t.context.copilotContext.mergeDocStatus( + workspace.id, + [] + ); + t.deepEqual(emptyResult, []); + + const basicDocs = [ + createDoc('doc1'), + createDoc('doc2'), + createDoc('doc3', 'failed'), + createDoc('doc4', 'processing'), + ]; + const basicResult = await t.context.copilotContext.mergeDocStatus( + workspace.id, + basicDocs + ); + t.snapshot( + basicResult.map(d => ({ id: d.id, status: d.status })), + 'basic doc status merge' + ); + + { + await createDocWithEmbedding('doc5'); + + const mixedDocs = [ + createDoc('doc5'), + createDoc('doc5', 'processing'), + createDoc('doc6'), + createDoc('doc6', 'failed'), + createDoc('doc7'), + ]; + const mixedResult = await t.context.copilotContext.mergeDocStatus( + workspace.id, + mixedDocs + ); + t.snapshot( + mixedResult.map(d => ({ id: d.id, status: d.status })), + 'mixed doc status merge' + ); + + const hasEmbeddingStub = Sinon.stub( + t.context.copilotContext, + 'hasWorkspaceEmbedding' + ).resolves(new Set()); + + const stubResult = await t.context.copilotContext.mergeDocStatus( + workspace.id, + [createDoc('doc5')] + ); + t.is(stubResult[0].status, ContextEmbedStatus.processing); + + hasEmbeddingStub.restore(); + } + + { + const testCases = [ + { + workspaceId: 'invalid-workspace', + docs: [{ id: 'doc1', createdAt: Date.now() }], + }, + { + workspaceId: workspace.id, + docs: [{ id: 'doc1', createdAt: Date.now(), status: undefined as any }], + }, + { + workspaceId: workspace.id, + docs: Array.from({ length: 100 }, (_, i) => ({ + id: `doc-${i}`, + createdAt: Date.now() + i, + })), + }, + ]; + + const results = await Promise.all( + testCases.map(testCase => + t.context.copilotContext.mergeDocStatus( + testCase.workspaceId, + testCase.docs + ) + ) + ); + + t.snapshot( + results.map((result, index) => ({ + case: index, + length: result.length, + statuses: result.map(d => d.status), + })), + 'edge cases results' + ); + } +}); + +test('should handle concurrent mergeDocStatus calls', async t => { + await t.context.db.snapshot.create({ + data: { + workspaceId: workspace.id, + id: 'concurrent-doc', + blob: Buffer.from([1, 1]), + state: Buffer.from([1, 1]), + updatedAt: new Date(), + createdAt: new Date(), + }, + }); + + await t.context.copilotContext.insertWorkspaceEmbedding( + workspace.id, + 'concurrent-doc', + [ + { + index: 0, + content: 'content', + embedding: Array.from({ length: 1024 }, () => 1), + }, + ] + ); + + const concurrentDocs = [ + [{ id: 'concurrent-doc', createdAt: Date.now() }], + [{ id: 'concurrent-doc', createdAt: Date.now() + 1000 }], + [{ id: 'non-existent-doc', createdAt: Date.now() }], + ]; + + const results = await Promise.all( + concurrentDocs.map(docs => + t.context.copilotContext.mergeDocStatus(workspace.id, docs) + ) + ); + + t.snapshot( + results.map((result, index) => ({ + call: index + 1, + status: result[0].status, + })), + 'concurrent calls results' + ); +}); diff --git a/packages/backend/server/src/models/copilot-context.ts b/packages/backend/server/src/models/copilot-context.ts index f346f9679b..ef9213afc4 100644 --- a/packages/backend/server/src/models/copilot-context.ts +++ b/packages/backend/server/src/models/copilot-context.ts @@ -91,7 +91,9 @@ export class CopilotContextModel extends BaseModel { const status = finishedDoc.has(doc.id) ? ContextEmbedStatus.finished : undefined; - doc.status = status || doc.status; + // NOTE: when the document has not been synchronized to the server or is in the embedding queue + // the status will be empty, fallback to processing if no status is provided + doc.status = status || doc.status || ContextEmbedStatus.processing; } return docs;