feat: improve context api (#10812)

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

@@ -1,4 +1,7 @@
 export { CopilotContextDocJob } from './job';
 export { CopilotContextResolver, CopilotContextRootResolver } from './resolver';
 export { CopilotContextService } from './service';
-export { type ContextFile, ContextFileStatus } from './types';
+export {
+  type ContextFile,
+  ContextEmbedStatus as ContextFileStatus,
+} from './types';

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

@@ -41,8 +41,8 @@ import { CopilotContextDocJob } from './job';
 import { CopilotContextService } from './service';
 import {
   ContextDoc,
+  ContextEmbedStatus,
   type ContextFile,
-  ContextFileStatus,
   DocChunkSimilarity,
   FileChunkSimilarity,
   MAX_EMBEDDABLE_SIZE,
@@ -94,13 +94,16 @@ export class CopilotContextType {
   workspaceId!: string;
 }
 
-registerEnumType(ContextFileStatus, { name: 'ContextFileStatus' });
+registerEnumType(ContextEmbedStatus, { name: 'ContextEmbedStatus' });
 
 @ObjectType()
 class CopilotContextDoc implements ContextDoc {
   @Field(() => ID)
   id!: string;
 
+  @Field(() => ContextEmbedStatus, { nullable: true })
+  status!: ContextEmbedStatus | null;
+
   @Field(() => SafeIntResolver)
   createdAt!: number;
 }
@@ -116,8 +119,8 @@ class CopilotContextFile implements ContextFile {
   @Field(() => SafeIntResolver)
   chunkSize!: number;
 
-  @Field(() => ContextFileStatus)
+  @Field(() => ContextEmbedStatus)
-  status!: ContextFileStatus;
+  status!: ContextEmbedStatus;
 
   @Field(() => String, { nullable: true })
   error!: string | null;

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

@@ -16,8 +16,8 @@ import { ContextSession } from './session';
 import {
   ContextConfig,
   ContextConfigSchema,
+  ContextEmbedStatus,
   ContextFile,
-  ContextFileStatus,
   EmbeddingClient,
 } from './types';
 import { checkEmbeddingAvailable } from './utils';
@@ -173,7 +173,7 @@ export class CopilotContextService implements OnModuleInit {
     await context.saveFileRecord(fileId, file => ({
       ...(file as ContextFile),
       chunkSize,
-      status: ContextFileStatus.finished,
+      status: ContextEmbedStatus.finished,
     }));
   }
 
@@ -187,7 +187,7 @@ export class CopilotContextService implements OnModuleInit {
     await context.saveFileRecord(fileId, file => ({
       ...(file as ContextFile),
       error,
-      status: ContextFileStatus.failed,
+      status: ContextEmbedStatus.failed,
     }));
   }
 }

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

@@ -6,8 +6,8 @@ import {
   ChunkSimilarity,
   ContextConfig,
   ContextDoc,
+  ContextEmbedStatus,
   ContextFile,
-  ContextFileStatus,
   ContextList,
   DocChunkSimilarity,
   EmbeddingClient,
@@ -54,7 +54,7 @@ export class ContextSession implements AsyncDisposable {
     if (doc) {
       return doc;
     }
-    const record = { id: docId, createdAt: Date.now() };
+    const record = { id: docId, createdAt: Date.now(), status: null };
     this.config.docs.push(record);
     await this.save();
     return record;
@@ -76,7 +76,7 @@ export class ContextSession implements AsyncDisposable {
     if (existsBlob) {
       // use exists file id if the blob exists
       // we assume that the file content pointed to by the same blobId is consistent.
-      if (existsBlob.status === ContextFileStatus.finished) {
+      if (existsBlob.status === ContextEmbedStatus.finished) {
         return existsBlob;
       }
       fileId = existsBlob.id;
@@ -179,7 +179,7 @@ export class ContextSession implements AsyncDisposable {
     if (file) {
       Object.assign(file, cb({ ...file }));
     } else {
-      const file = { id: fileId, status: ContextFileStatus.processing };
+      const file = { id: fileId, status: ContextEmbedStatus.processing };
       files.push(cb(file));
     }
     await this.save(tx);

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

@@ -26,7 +26,7 @@ declare global {
 
 export const MAX_EMBEDDABLE_SIZE = 50 * OneMB;
 
-export enum ContextFileStatus {
+export enum ContextEmbedStatus {
   processing = 'processing',
   finished = 'finished',
   failed = 'failed',
@@ -40,9 +40,9 @@ export const ContextConfigSchema = z.object({
       chunkSize: z.number(),
       name: z.string(),
       status: z.enum([
-        ContextFileStatus.processing,
+        ContextEmbedStatus.processing,
-        ContextFileStatus.finished,
+        ContextEmbedStatus.finished,
-        ContextFileStatus.failed,
+        ContextEmbedStatus.failed,
       ]),
       error: z.string().nullable(),
       blobId: z.string(),
@@ -52,6 +52,15 @@ export const ContextConfigSchema = z.object({
   docs: z
     .object({
       id: z.string(),
+      // status for workspace doc embedding progress
+      // only exists when the client submits the doc embedding task
+      status: z
+        .enum([
+          ContextEmbedStatus.processing,
+          ContextEmbedStatus.finished,
+          ContextEmbedStatus.failed,
+        ])
+        .nullable(),
       createdAt: z.number(),
     })
     .array(),

packages/backend/server/src/schema.gql

@@ -35,7 +35,7 @@ type ChatMessage {
   role: String!
 }
 
-enum ContextFileStatus {
+enum ContextEmbedStatus {
   failed
   finished
   processing
@@ -95,6 +95,7 @@ type CopilotContext {
 type CopilotContextDoc {
   createdAt: SafeInt!
   id: ID!
+  status: ContextEmbedStatus
 }
 
 type CopilotContextFile {
@@ -104,7 +105,7 @@ type CopilotContextFile {
   error: String
   id: ID!
   name: String!
-  status: ContextFileStatus!
+  status: ContextEmbedStatus!
 }
 
 type CopilotContextFileNotSupportedDataType {

packages/frontend/graphql/src/graphql/copilot-context-doc-add.gql

@@ -2,5 +2,6 @@ mutation addContextDoc($options: AddContextDocInput!) {
   addContextDoc(options: $options) {
     id
     createdAt
+    status
   }
 }

packages/frontend/graphql/src/graphql/copilot-context-list-docs-and-files.gql

@@ -4,6 +4,7 @@ query listContextDocsAndFiles($workspaceId: String!, $sessionId: String!, $conte
       contexts(sessionId: $sessionId, contextId: $contextId) {
         docs {
           id
+          status
           createdAt
         }
         files {

packages/frontend/graphql/src/graphql/index.ts

@@ -149,6 +149,7 @@ export const addContextDocMutation = {
   addContextDoc(options: $options) {
     id
     createdAt
+    status
   }
 }`,
 };
@@ -240,6 +241,7 @@ export const listContextDocsAndFilesQuery = {
       contexts(sessionId: $sessionId, contextId: $contextId) {
         docs {
           id
+          status
           createdAt
         }
         files {

packages/frontend/graphql/src/schema.ts

@@ -73,7 +73,7 @@ export interface ChatMessage {
   role: Scalars['String']['output'];
 }
 
-export enum ContextFileStatus {
+export enum ContextEmbedStatus {
   failed = 'failed',
   finished = 'finished',
   processing = 'processing',
@@ -167,6 +167,7 @@ export interface CopilotContextDoc {
   __typename?: 'CopilotContextDoc';
   createdAt: Scalars['SafeInt']['output'];
   id: Scalars['ID']['output'];
+  status: Maybe<ContextEmbedStatus>;
 }
 
 export interface CopilotContextFile {
@@ -177,7 +178,7 @@ export interface CopilotContextFile {
   error: Maybe<Scalars['String']['output']>;
   id: Scalars['ID']['output'];
   name: Scalars['String']['output'];
-  status: ContextFileStatus;
+  status: ContextEmbedStatus;
 }
 
 export interface CopilotContextFileNotSupportedDataType {
@@ -2330,6 +2331,7 @@ export type AddContextDocMutation = {
     __typename?: 'CopilotContextDoc';
     id: string;
     createdAt: number;
+    status: ContextEmbedStatus | null;
   };
 };
 
@@ -2356,7 +2358,7 @@ export type AddContextFileMutation = {
     name: string;
     chunkSize: number;
     error: string | null;
-    status: ContextFileStatus;
+    status: ContextEmbedStatus;
     blobId: string;
   };
 };
@@ -2387,7 +2389,7 @@ export type ListContextFilesQuery = {
           blobId: string;
           chunkSize: number;
           error: string | null;
-          status: ContextFileStatus;
+          status: ContextEmbedStatus;
           createdAt: number;
         }>;
       }>;
@@ -2447,6 +2449,7 @@ export type ListContextDocsAndFilesQuery = {
         docs: Array<{
           __typename?: 'CopilotContextDoc';
           id: string;
+          status: ContextEmbedStatus | null;
           createdAt: number;
         }>;
         files: Array<{
@@ -2456,7 +2459,7 @@ export type ListContextDocsAndFilesQuery = {
           blobId: string;
           chunkSize: number;
           error: string | null;
-          status: ContextFileStatus;
+          status: ContextEmbedStatus;
           createdAt: number;
         }>;
       }>;