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:
darkskygit
2025-06-10 02:35:09 +00:00
parent e82c9d2ddc
commit c846c57a12
6 changed files with 368 additions and 1 deletions

View File

@@ -52,6 +52,7 @@ Generated by [AVA](https://avajs.dev).
[
{
id: 'docId1',
status: 'processing',
},
]

View File

@@ -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',
},
]

View File

@@ -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'
);
});

View File

@@ -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;