mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-02-12 20:38:52 +00:00
fix(server): list context status (#12771)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved handling of document statuses to ensure documents without a finished or existing status are now explicitly marked as "processing" instead of remaining undefined. - **Tests** - Added comprehensive new tests and snapshot entries to verify document status merging, including edge cases and concurrent operations, ensuring robust and consistent behavior. - **Enhancements** - Updated context document listings to display the processing status for relevant documents. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
@@ -52,6 +52,7 @@ Generated by [AVA](https://avajs.dev).
|
||||
[
|
||||
{
|
||||
id: 'docId1',
|
||||
status: 'processing',
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
Binary file not shown.
@@ -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',
|
||||
},
|
||||
]
|
||||
|
||||
Binary file not shown.
@@ -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<string>());
|
||||
|
||||
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'
|
||||
);
|
||||
});
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user