fix(server): make status optional (#11176)

packages/backend/server/src/__tests__/models/copilot-context.spec.ts

@@ -84,7 +84,6 @@ test('should update context', async t => {
   const doc = {
     id: docId,
     createdAt: Date.now(),
-    status: null,
   };
   config?.docs.push(doc);
   await t.context.copilotContext.update(contextId, { config });

packages/backend/server/src/models/common/copilot.ts

@@ -31,27 +31,22 @@ export enum ContextCategories {
   Collection = 'collection',
 }
 
-export const ContextDocSchema = z.object({
+const ContextEmbedStatusSchema = z.enum([
+  ContextEmbedStatus.processing,
+  ContextEmbedStatus.finished,
+  ContextEmbedStatus.failed,
+]);
+
+const ContextDocSchema = z.object({
   id: z.string(),
   createdAt: z.number(),
-  status: z
-    .enum([
-      ContextEmbedStatus.processing,
-      ContextEmbedStatus.finished,
-      ContextEmbedStatus.failed,
-    ])
-    .nullable(),
 });
 
 export const ContextFileSchema = z.object({
   id: z.string(),
   chunkSize: z.number(),
   name: z.string(),
-  status: z.enum([
-    ContextEmbedStatus.processing,
-    ContextEmbedStatus.finished,
-    ContextEmbedStatus.failed,
-  ]),
+  status: ContextEmbedStatusSchema,
   error: z.string().nullable(),
   blobId: z.string(),
   createdAt: z.number(),
@@ -60,14 +55,18 @@ export const ContextFileSchema = z.object({
 export const ContextCategorySchema = z.object({
   id: z.string(),
   type: z.enum([ContextCategories.Tag, ContextCategories.Collection]),
-  docs: ContextDocSchema.array(),
+  docs: ContextDocSchema.merge(
+    z.object({ status: ContextEmbedStatusSchema })
+  ).array(),
   createdAt: z.number(),
 });
 
 export const ContextConfigSchema = z.object({
   workspaceId: z.string(),
   files: ContextFileSchema.array(),
-  docs: ContextDocSchema.array(),
+  docs: ContextDocSchema.merge(
+    z.object({ status: ContextEmbedStatusSchema.optional() })
+  ).array(),
   categories: ContextCategorySchema.array(),
 });
 
@@ -76,9 +75,9 @@ export const MinimalContextConfigSchema = ContextConfigSchema.pick({
 });
 
 export type ContextCategory = z.infer<typeof ContextCategorySchema>;
-export type ContextDoc = z.infer<typeof ContextDocSchema>;
-export type ContextFile = z.infer<typeof ContextFileSchema>;
 export type ContextConfig = z.infer<typeof ContextConfigSchema>;
+export type ContextDoc = z.infer<typeof ContextConfigSchema>['docs'][number];
+export type ContextFile = z.infer<typeof ContextConfigSchema>['files'][number];
 export type ContextListItem = ContextDoc | ContextFile;
 export type ContextList = ContextListItem[];
 

packages/backend/server/src/models/copilot-context.ts

@@ -83,18 +83,15 @@ export class CopilotContextModel extends BaseModel {
     return row;
   }
 
-  async mergeDocStatus(
-    workspaceId: string,
-    docs: (ContextDoc & { status?: ContextEmbedStatus | null })[]
-  ) {
+  async mergeDocStatus(workspaceId: string, docs: ContextDoc[]) {
     const docIds = Array.from(new Set(docs.map(doc => doc.id)));
     const finishedDoc = await this.hasWorkspaceEmbedding(workspaceId, docIds);
 
     for (const doc of docs) {
       const status = finishedDoc.has(doc.id)
         ? ContextEmbedStatus.finished
-        : null;
-      doc.status = status;
+        : undefined;
+      doc.status = status || doc.status;
     }
 
     return docs;

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

@@ -127,7 +127,7 @@ export class CopilotContextType {
 registerEnumType(ContextCategories, { name: 'ContextCategories' });
 
 @ObjectType()
-class CopilotDocType implements ContextDoc {
+class CopilotDocType implements Omit<ContextDoc, 'status'> {
   @Field(() => ID)
   id!: string;
 
@@ -418,12 +418,12 @@ export class CopilotContextResolver {
     description: 'list files in context',
   })
   @CallMetric('ai', 'context_file_list')
-  async docs(@Parent() context: CopilotContextType): Promise<ContextDoc[]> {
+  async docs(@Parent() context: CopilotContextType): Promise<CopilotDocType[]> {
     const session = await this.context.get(context.id);
     const docs = session.docs;
     await this.models.copilotContext.mergeDocStatus(session.workspaceId, docs);
 
-    return docs;
+    return docs.map(doc => ({ ...doc, status: doc.status || null }));
   }
 
   @ResolveField(() => [CopilotContextFile], {
@@ -512,11 +512,11 @@ export class CopilotContextResolver {
   async addContextDoc(
     @Args({ name: 'options', type: () => AddContextDocInput })
     options: AddContextDocInput
-  ) {
+  ): Promise<CopilotDocType> {
     const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
     await using lock = await this.mutex.acquire(lockFlag);
     if (!lock) {
-      return new TooManyRequest('Server is busy');
+      throw new TooManyRequest('Server is busy');
     }
     const session = await this.context.get(options.contextId);
 
@@ -530,7 +530,7 @@ export class CopilotContextResolver {
         },
       ]);
 
-      return record;
+      return { ...record, status: record.status || null };
     } catch (e: any) {
       throw new CopilotFailedToModifyContext({
         contextId: options.contextId,

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

@@ -114,7 +114,7 @@ export class ContextSession implements AsyncDisposable {
     if (doc) {
       return doc;
     }
-    const record = { id: docId, createdAt: Date.now(), status: null };
+    const record = { id: docId, createdAt: Date.now() };
     this.config.docs.push(record);
     await this.save();
     return record;