mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-02-15 21:41:52 +08: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',
|
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
|
> should return true when embedding table is available
|
||||||
|
|
||||||
true
|
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 { AiSession, PrismaClient, User, Workspace } from '@prisma/client';
|
||||||
import ava, { TestFn } from 'ava';
|
import ava, { TestFn } from 'ava';
|
||||||
|
import Sinon from 'sinon';
|
||||||
|
|
||||||
import { Config } from '../../base';
|
import { Config } from '../../base';
|
||||||
|
import { ContextEmbedStatus } from '../../models/common/copilot';
|
||||||
import { CopilotContextModel } from '../../models/copilot-context';
|
import { CopilotContextModel } from '../../models/copilot-context';
|
||||||
import { CopilotSessionModel } from '../../models/copilot-session';
|
import { CopilotSessionModel } from '../../models/copilot-session';
|
||||||
import { CopilotWorkspaceConfigModel } from '../../models/copilot-workspace';
|
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');
|
// 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)
|
const status = finishedDoc.has(doc.id)
|
||||||
? ContextEmbedStatus.finished
|
? ContextEmbedStatus.finished
|
||||||
: undefined;
|
: 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;
|
return docs;
|
||||||
|
|||||||
Reference in New Issue
Block a user