feat(server): basic context api (#10056)

fix CLOUD-97
fix CLOUD-98

packages/backend/server/migrations/20250210083416_ai_context/migration.sql

@@ -0,0 +1,13 @@
+-- CreateTable
+CREATE TABLE "ai_contexts" (
+    "id" VARCHAR NOT NULL,
+    "session_id" VARCHAR NOT NULL,
+    "config" JSON NOT NULL,
+    "created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMPTZ(3) NOT NULL,
+
+    CONSTRAINT "ai_contexts_pkey" PRIMARY KEY ("id")
+);
+
+-- AddForeignKey
+ALTER TABLE "ai_contexts" ADD CONSTRAINT "ai_contexts_session_id_fkey" FOREIGN KEY ("session_id") REFERENCES "ai_sessions_metadata"("id") ON DELETE CASCADE ON UPDATE CASCADE;

packages/backend/server/package.json

@@ -82,7 +82,7 @@
     "nestjs-cls": "^5.0.0",
     "nodemailer": "^6.9.16",
     "on-headers": "^1.0.2",
-    "openai": "^4.76.2",
+    "openai": "^4.83.0",
     "piscina": "^5.0.0-alpha.0",
     "prisma": "^5.22.0",
     "react": "19.0.0",

packages/backend/server/schema.prisma

@@ -402,12 +402,26 @@ model AiSession {
   user     User               @relation(fields: [userId], references: [id], onDelete: Cascade)
   prompt   AiPrompt           @relation(fields: [promptName], references: [name], onDelete: Cascade)
   messages AiSessionMessage[]
+  context  AiContext[]
 
   @@index([userId])
   @@index([userId, workspaceId])
   @@map("ai_sessions_metadata")
 }
 
+model AiContext {
+  id        String @id @default(uuid()) @db.VarChar
+  sessionId String @map("session_id") @db.VarChar
+  config    Json   @db.Json
+
+  createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
+  updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3)
+
+  session AiSession @relation(fields: [sessionId], references: [id], onDelete: Cascade)
+
+  @@map("ai_contexts")
+}
+
 model DataMigration {
   id         String    @id @default(uuid()) @db.VarChar
   name       String    @unique @db.VarChar

packages/backend/server/src/__tests__/__snapshots__/copilot.e2e.ts.md

@@ -0,0 +1,26 @@
+# Snapshot report for `src/__tests__/copilot.e2e.ts`
+
+The actual snapshot is saved in `copilot.e2e.ts.snap`.
+
+Generated by [AVA](https://avajs.dev).
+
+## should be able to manage context
+
+> should list context files
+
+    [
+      {
+        id: 'docId1',
+      },
+    ]
+
+> should list context docs
+
+    [
+      {
+        blobId: 'fileId1',
+        chunkSize: 3,
+        name: 'sample.pdf',
+        status: 'finished',
+      },
+    ]

packages/backend/server/src/__tests__/copilot.e2e.ts

@@ -8,6 +8,7 @@ import { ConfigModule } from '../base/config';
 import { AuthService } from '../core/auth';
 import { WorkspaceModule } from '../core/workspaces';
 import { CopilotModule } from '../plugins/copilot';
+import { CopilotContextService } from '../plugins/copilot/context';
 import { prompts, PromptService } from '../plugins/copilot/prompt';
 import {
   CopilotProviderService,
@@ -27,15 +28,19 @@ import {
   TestUser,
 } from './utils';
 import {
+  addContextDoc,
   array2sse,
   chatWithImages,
   chatWithText,
   chatWithTextStream,
   chatWithWorkflow,
+  createCopilotContext,
   createCopilotMessage,
   createCopilotSession,
   forkCopilotSession,
   getHistories,
+  listContext,
+  listContextFiles,
   MockCopilotTestProvider,
   sse2array,
   textToEventStream,
@@ -46,6 +51,7 @@ import {
 const test = ava as TestFn<{
   auth: AuthService;
   app: TestingApp;
+  context: CopilotContextService;
   prompt: PromptService;
   provider: CopilotProviderService;
   storage: CopilotStorage;
@@ -77,11 +83,13 @@ test.before(async t => {
   });
 
   const auth = app.get(AuthService);
+  const context = app.get(CopilotContextService);
   const prompt = app.get(PromptService);
   const storage = app.get(CopilotStorage);
 
   t.context.app = app;
   t.context.auth = auth;
+  t.context.context = context;
   t.context.prompt = prompt;
   t.context.storage = storage;
 });
@@ -678,3 +686,46 @@ test('should be able to search image from unsplash', async t => {
   const resp = await unsplashSearch(app);
   t.not(resp.status, 404, 'route should be exists');
 });
+
+test('should be able to manage context', async t => {
+  const { app } = t.context;
+
+  const { id: workspaceId } = await createWorkspace(app);
+  const sessionId = await createCopilotSession(
+    app,
+    workspaceId,
+    randomUUID(),
+    promptName
+  );
+
+  {
+    await t.throwsAsync(
+      createCopilotContext(app, workspaceId, randomUUID()),
+      { instanceOf: Error },
+      'should throw error if create context with invalid session id'
+    );
+
+    const context = createCopilotContext(app, workspaceId, sessionId);
+    await t.notThrowsAsync(context, 'should create context with chat session');
+
+    const list = await listContext(app, workspaceId, sessionId);
+    t.deepEqual(
+      list.map(f => ({ id: f.id })),
+      [{ id: await context }],
+      'should list context'
+    );
+  }
+
+  {
+    const contextId = await createCopilotContext(app, workspaceId, sessionId);
+
+    await addContextDoc(app, contextId, 'docId1');
+
+    const { docs } =
+      (await listContextFiles(app, workspaceId, sessionId, contextId)) || {};
+    t.snapshot(
+      docs?.map(({ createdAt: _, ...d }) => d),
+      'should list context files'
+    );
+  }
+});

packages/backend/server/src/__tests__/copilot.spec.ts

@@ -1,3 +1,5 @@
+import { randomUUID } from 'node:crypto';
+
 import type { TestFn } from 'ava';
 import ava from 'ava';
 import Sinon from 'sinon';
@@ -6,6 +8,7 @@ import { ConfigModule } from '../base/config';
 import { AuthService } from '../core/auth';
 import { QuotaModule } from '../core/quota';
 import { CopilotModule } from '../plugins/copilot';
+import { CopilotContextService } from '../plugins/copilot/context';
 import { prompts, PromptService } from '../plugins/copilot/prompt';
 import {
   CopilotProviderService,
@@ -44,6 +47,7 @@ import { MockCopilotTestProvider, WorkflowTestCases } from './utils/copilot';
 const test = ava as TestFn<{
   auth: AuthService;
   module: TestingModule;
+  context: CopilotContextService;
   prompt: PromptService;
   provider: CopilotProviderService;
   session: ChatSessionService;
@@ -81,6 +85,7 @@ test.before(async t => {
   });
 
   const auth = module.get(AuthService);
+  const context = module.get(CopilotContextService);
   const prompt = module.get(PromptService);
   const provider = module.get(CopilotProviderService);
   const session = module.get(ChatSessionService);
@@ -88,6 +93,7 @@ test.before(async t => {
 
   t.context.module = module;
   t.context.auth = auth;
+  t.context.context = context;
   t.context.prompt = prompt;
   t.context.provider = provider;
   t.context.session = session;
@@ -1247,3 +1253,52 @@ test('CitationParser should not replace chunks of citation already with URLs', t
   ].join('\n');
   t.is(result, expected);
 });
+
+// ==================== context ====================
+test('should be able to manage context', async t => {
+  const { context, prompt, session } = t.context;
+
+  await prompt.set('prompt', 'model', [
+    { role: 'system', content: 'hello {{word}}' },
+  ]);
+  const chatSession = await session.create({
+    docId: 'test',
+    workspaceId: 'test',
+    userId,
+    promptName: 'prompt',
+  });
+
+  {
+    await t.throwsAsync(
+      context.create(randomUUID()),
+      { instanceOf: Error },
+      'should throw error if create context with invalid session id'
+    );
+
+    const session = context.create(chatSession);
+    await t.notThrowsAsync(session, 'should create context with chat session');
+
+    await t.notThrowsAsync(
+      context.get((await session).id),
+      'should get context after create'
+    );
+
+    await t.throwsAsync(
+      context.get(randomUUID()),
+      { instanceOf: Error },
+      'should throw error if get context with invalid id'
+    );
+  }
+
+  {
+    const session = await context.create(chatSession);
+
+    const docId = randomUUID();
+    await session.addDocRecord(docId);
+    const docs = session.listDocs().map(d => d.id);
+    t.deepEqual(docs, [docId], 'should list doc id');
+
+    await session.removeDocRecord(docId);
+    t.deepEqual(session.listDocs(), [], 'should remove doc id');
+  }
+});

packages/backend/server/src/__tests__/utils/copilot.ts

@@ -23,6 +23,7 @@ import {
   WorkflowNodeType,
   WorkflowParams,
 } from '../../plugins/copilot/workflow/types';
+import { gql } from './common';
 import { TestingApp } from './testing-app';
 import { sleep } from './utils';
 
@@ -209,6 +210,216 @@ export async function forkCopilotSession(
   return res.forkCopilotSession;
 }
 
+export async function createCopilotContext(
+  app: TestingApp,
+  workspaceId: string,
+  sessionId: string
+): Promise<string> {
+  const res = await app.gql(`
+        mutation {
+          createCopilotContext(workspaceId: "${workspaceId}", sessionId: "${sessionId}")
+        }
+      `);
+
+  return res.createCopilotContext;
+}
+
+export async function matchContext(
+  app: TestingApp,
+  contextId: string,
+  content: string,
+  limit: number
+): Promise<
+  | {
+      fileId: string;
+      chunk: number;
+      content: string;
+      distance: number | null;
+    }[]
+  | undefined
+> {
+  const res = await app.gql(
+    `
+        mutation matchContext($content: String!, $contextId: String!, $limit: SafeInt) {
+          matchContext(content: $content, contextId: $contextId, limit: $limit) {
+            fileId
+            chunk
+            content
+            distance
+          }
+        }
+      `,
+    { contextId, content, limit }
+  );
+
+  return res.matchContext;
+}
+
+export async function listContext(
+  app: TestingApp,
+  workspaceId: string,
+  sessionId: string
+): Promise<
+  {
+    id: string;
+    workspaceId: string;
+  }[]
+> {
+  const res = await app.gql(`
+        query {
+          currentUser {
+            copilot(workspaceId: "${workspaceId}") {
+              contexts(sessionId: "${sessionId}") {
+                id
+                workspaceId
+              }
+            }
+          }
+        }
+      `);
+
+  return res.currentUser?.copilot?.contexts;
+}
+
+export async function addContextFile(
+  app: TestingApp,
+  contextId: string,
+  blobId: string,
+  fileName: string,
+  content: Buffer
+): Promise<{ id: string }[]> {
+  const res = await app
+    .POST(gql)
+    .set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
+    .field(
+      'operations',
+      JSON.stringify({
+        query: `
+          mutation addContextFile($options: AddContextFileInput!, $content: Upload!) {
+            addContextFile(content: $content, options: $options) {
+              id
+            }
+          }
+        `,
+        variables: {
+          content: null,
+          options: { contextId, blobId, fileName },
+        },
+      })
+    )
+    .field('map', JSON.stringify({ '0': ['variables.content'] }))
+    .attach('0', content, {
+      filename: fileName,
+      contentType: 'application/octet-stream',
+    })
+    .expect(200);
+
+  return res.body.data.addContextFile;
+}
+
+export async function removeContextFile(
+  app: TestingApp,
+  contextId: string,
+  fileId: string
+): Promise<string> {
+  const res = await app.gql(
+    `
+        mutation removeContextFile($options: RemoveContextFileInput!) {
+          removeContextFile(options: $options)
+        }
+    `,
+    { options: { contextId, fileId } }
+  );
+
+  return res.removeContextFile;
+}
+
+export async function addContextDoc(
+  app: TestingApp,
+  contextId: string,
+  docId: string
+): Promise<{ id: string }[]> {
+  const res = await app.gql(
+    `
+          mutation addContextDoc($options: AddContextDocInput!) {
+            addContextDoc(options: $options) {
+              id
+            }
+          }
+        `,
+    { options: { contextId, docId } }
+  );
+
+  return res.addContextDoc;
+}
+
+export async function removeContextDoc(
+  app: TestingApp,
+  contextId: string,
+  docId: string
+): Promise<string> {
+  const res = await app.gql(
+    `
+      mutation removeContextDoc($options: RemoveContextFileInput!) {
+        removeContextDoc(options: $options)
+      }
+    `,
+    { options: { contextId, docId } }
+  );
+
+  return res.removeContextDoc;
+}
+
+export async function listContextFiles(
+  app: TestingApp,
+  workspaceId: string,
+  sessionId: string,
+  contextId: string
+): Promise<
+  | {
+      docs: {
+        id: string;
+        createdAt: number;
+      }[];
+      files: {
+        id: string;
+        name: string;
+        blobId: string;
+        chunkSize: number;
+        status: string;
+        createdAt: number;
+      }[];
+    }
+  | undefined
+> {
+  const res = await app.gql(`
+        query {
+          currentUser {
+            copilot(workspaceId: "${workspaceId}") {
+              contexts(sessionId: "${sessionId}", contextId: "${contextId}") {
+                docs {
+                  id
+                  createdAt
+                }
+                files {
+                  id
+                  name
+                  blobId
+                  chunkSize
+                  status
+                  createdAt
+                }
+              }
+            }
+          }
+        }
+      `);
+
+  const { docs, files } = res.currentUser?.copilot?.contexts?.[0] || {};
+
+  return { docs, files };
+}
+
 export async function createCopilotMessage(
   app: TestingApp,
   sessionId: string,

packages/backend/server/src/base/error/def.ts

@@ -1,4 +1,5 @@
 import { STATUS_CODES } from 'node:http';
+import { escape } from 'node:querystring';
 
 import { HttpStatus, Logger } from '@nestjs/common';
 import { ClsServiceManager } from 'nestjs-cls';
@@ -605,6 +606,29 @@ export const USER_FRIENDLY_ERRORS = {
     message: ({ provider, kind, message }) =>
       `Provider ${provider} failed with ${kind} error: ${message || 'unknown'}`,
   },
+  copilot_invalid_context: {
+    type: 'invalid_input',
+    args: { contextId: 'string' },
+    message: ({ contextId }) => `Invalid copilot context ${contextId}.`,
+  },
+  copilot_context_file_not_supported: {
+    type: 'bad_request',
+    args: { fileName: 'string', message: 'string' },
+    message: ({ fileName, message }) =>
+      `File ${fileName} is not supported to use as context: ${message}`,
+  },
+  copilot_failed_to_modify_context: {
+    type: 'internal_server_error',
+    args: { contextId: 'string', message: 'string' },
+    message: ({ contextId, message }) =>
+      `Failed to modify context ${contextId}: ${message}`,
+  },
+  copilot_failed_to_match_context: {
+    type: 'internal_server_error',
+    args: { contextId: 'string', content: 'string', message: 'string' },
+    message: ({ contextId, content, message }) =>
+      `Failed to match context ${contextId} with "${escape(content)}": ${message}`,
+  },
 
   // Quota & Limit errors
   blob_quota_exceeded: {

packages/backend/server/src/base/error/errors.gen.ts

@@ -608,6 +608,50 @@ export class CopilotProviderSideError extends UserFriendlyError {
     super('internal_server_error', 'copilot_provider_side_error', message, args);
   }
 }
+@ObjectType()
+class CopilotInvalidContextDataType {
+  @Field() contextId!: string
+}
+
+export class CopilotInvalidContext extends UserFriendlyError {
+  constructor(args: CopilotInvalidContextDataType, message?: string | ((args: CopilotInvalidContextDataType) => string)) {
+    super('invalid_input', 'copilot_invalid_context', message, args);
+  }
+}
+@ObjectType()
+class CopilotContextFileNotSupportedDataType {
+  @Field() fileName!: string
+  @Field() message!: string
+}
+
+export class CopilotContextFileNotSupported extends UserFriendlyError {
+  constructor(args: CopilotContextFileNotSupportedDataType, message?: string | ((args: CopilotContextFileNotSupportedDataType) => string)) {
+    super('bad_request', 'copilot_context_file_not_supported', message, args);
+  }
+}
+@ObjectType()
+class CopilotFailedToModifyContextDataType {
+  @Field() contextId!: string
+  @Field() message!: string
+}
+
+export class CopilotFailedToModifyContext extends UserFriendlyError {
+  constructor(args: CopilotFailedToModifyContextDataType, message?: string | ((args: CopilotFailedToModifyContextDataType) => string)) {
+    super('internal_server_error', 'copilot_failed_to_modify_context', message, args);
+  }
+}
+@ObjectType()
+class CopilotFailedToMatchContextDataType {
+  @Field() contextId!: string
+  @Field() content!: string
+  @Field() message!: string
+}
+
+export class CopilotFailedToMatchContext extends UserFriendlyError {
+  constructor(args: CopilotFailedToMatchContextDataType, message?: string | ((args: CopilotFailedToMatchContextDataType) => string)) {
+    super('internal_server_error', 'copilot_failed_to_match_context', message, args);
+  }
+}
 
 export class BlobQuotaExceeded extends UserFriendlyError {
   constructor(message?: string) {
@@ -801,6 +845,10 @@ export enum ErrorNames {
   COPILOT_PROMPT_NOT_FOUND,
   COPILOT_PROMPT_INVALID,
   COPILOT_PROVIDER_SIDE_ERROR,
+  COPILOT_INVALID_CONTEXT,
+  COPILOT_CONTEXT_FILE_NOT_SUPPORTED,
+  COPILOT_FAILED_TO_MODIFY_CONTEXT,
+  COPILOT_FAILED_TO_MATCH_CONTEXT,
   BLOB_QUOTA_EXCEEDED,
   MEMBER_QUOTA_EXCEEDED,
   COPILOT_QUOTA_EXCEEDED,
@@ -825,5 +873,5 @@ registerEnumType(ErrorNames, {
 export const ErrorDataUnionType = createUnionType({
   name: 'ErrorDataUnion',
   types: () =>
-    [QueryTooLongDataType, WrongSignInCredentialsDataType, UnknownOauthProviderDataType, MissingOauthQueryParameterDataType, InvalidEmailDataType, InvalidPasswordLengthDataType, WorkspacePermissionNotFoundDataType, SpaceNotFoundDataType, MemberNotFoundInSpaceDataType, NotInSpaceDataType, AlreadyInSpaceDataType, SpaceAccessDeniedDataType, SpaceOwnerNotFoundDataType, SpaceShouldHaveOnlyOneOwnerDataType, DocNotFoundDataType, DocAccessDeniedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, ExpectToGrantDocUserRolesDataType, ExpectToRevokeDocUserRolesDataType, ExpectToUpdateDocUserRoleDataType, UnsupportedSubscriptionPlanDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CopilotMessageNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderSideErrorDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType, InvalidLicenseUpdateParamsDataType, WorkspaceMembersExceedLimitToDowngradeDataType] as const,
+    [QueryTooLongDataType, WrongSignInCredentialsDataType, UnknownOauthProviderDataType, MissingOauthQueryParameterDataType, InvalidEmailDataType, InvalidPasswordLengthDataType, WorkspacePermissionNotFoundDataType, SpaceNotFoundDataType, MemberNotFoundInSpaceDataType, NotInSpaceDataType, AlreadyInSpaceDataType, SpaceAccessDeniedDataType, SpaceOwnerNotFoundDataType, SpaceShouldHaveOnlyOneOwnerDataType, DocNotFoundDataType, DocAccessDeniedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, ExpectToGrantDocUserRolesDataType, ExpectToRevokeDocUserRolesDataType, ExpectToUpdateDocUserRoleDataType, UnsupportedSubscriptionPlanDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CopilotMessageNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderSideErrorDataType, CopilotInvalidContextDataType, CopilotContextFileNotSupportedDataType, CopilotFailedToModifyContextDataType, CopilotFailedToMatchContextDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType, InvalidLicenseUpdateParamsDataType, WorkspaceMembersExceedLimitToDowngradeDataType] as const,
 });

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

@@ -0,0 +1,3 @@
+export { CopilotContextResolver, CopilotContextRootResolver } from './resolver';
+export { CopilotContextService } from './service';
+export { type ContextFile, ContextFileStatus } from './types';

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

@@ -0,0 +1,260 @@
+import {
+  Args,
+  Field,
+  ID,
+  InputType,
+  Mutation,
+  ObjectType,
+  Parent,
+  registerEnumType,
+  ResolveField,
+  Resolver,
+} from '@nestjs/graphql';
+import { SafeIntResolver } from 'graphql-scalars';
+
+import {
+  CallMetric,
+  CopilotFailedToModifyContext,
+  CopilotSessionNotFound,
+  RequestMutex,
+  Throttle,
+  TooManyRequest,
+} from '../../../base';
+import { CurrentUser } from '../../../core/auth';
+import { COPILOT_LOCKER, CopilotType } from '../resolver';
+import { ChatSessionService } from '../session';
+import { CopilotContextService } from './service';
+import { ContextDoc, type ContextFile, ContextFileStatus } from './types';
+
+@InputType()
+class AddContextDocInput {
+  @Field(() => String)
+  contextId!: string;
+
+  @Field(() => String)
+  docId!: string;
+}
+
+@InputType()
+class RemoveContextFileInput {
+  @Field(() => String)
+  contextId!: string;
+
+  @Field(() => String)
+  fileId!: string;
+}
+
+@ObjectType('CopilotContext')
+export class CopilotContextType {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => String)
+  workspaceId!: string;
+}
+
+registerEnumType(ContextFileStatus, { name: 'ContextFileStatus' });
+
+@ObjectType()
+class CopilotContextDoc implements ContextDoc {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => SafeIntResolver)
+  createdAt!: number;
+}
+
+@ObjectType()
+class CopilotContextFile implements ContextFile {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => String)
+  name!: string;
+
+  @Field(() => SafeIntResolver)
+  chunkSize!: number;
+
+  @Field(() => ContextFileStatus)
+  status!: ContextFileStatus;
+
+  @Field(() => String)
+  blobId!: string;
+
+  @Field(() => SafeIntResolver)
+  createdAt!: number;
+}
+
+@ObjectType()
+class CopilotContextListItem {
+  @Field(() => ID)
+  id!: string;
+
+  @Field(() => SafeIntResolver)
+  createdAt!: number;
+
+  @Field(() => String, { nullable: true })
+  name!: string;
+
+  @Field(() => SafeIntResolver, { nullable: true })
+  chunkSize!: number;
+
+  @Field(() => ContextFileStatus, { nullable: true })
+  status!: ContextFileStatus;
+
+  @Field(() => String, { nullable: true })
+  blobId!: string;
+}
+
+@Throttle()
+@Resolver(() => CopilotType)
+export class CopilotContextRootResolver {
+  constructor(
+    private readonly mutex: RequestMutex,
+    private readonly chatSession: ChatSessionService,
+    private readonly context: CopilotContextService
+  ) {}
+
+  private async checkChatSession(
+    user: CurrentUser,
+    sessionId: string,
+    workspaceId?: string
+  ): Promise<void> {
+    const session = await this.chatSession.get(sessionId);
+    if (
+      !session ||
+      session.config.workspaceId !== workspaceId ||
+      session.config.userId !== user.id
+    ) {
+      throw new CopilotSessionNotFound();
+    }
+  }
+
+  @ResolveField(() => [CopilotContextType], {
+    description: 'Get the context list of a session',
+    complexity: 2,
+  })
+  @CallMetric('ai', 'context_create')
+  async contexts(
+    @Parent() copilot: CopilotType,
+    @CurrentUser() user: CurrentUser,
+    @Args('sessionId') sessionId: string,
+    @Args('contextId', { nullable: true }) contextId?: string
+  ) {
+    const lockFlag = `${COPILOT_LOCKER}:context:${sessionId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      return new TooManyRequest('Server is busy');
+    }
+    await this.checkChatSession(user, sessionId, copilot.workspaceId);
+
+    if (contextId) {
+      const context = await this.context.get(contextId);
+      if (context) return [context];
+    } else {
+      const context = await this.context.getBySessionId(sessionId);
+      if (context) return [context];
+    }
+    return [];
+  }
+
+  @Mutation(() => String, {
+    description: 'Create a context session',
+  })
+  @CallMetric('ai', 'context_create')
+  async createCopilotContext(
+    @CurrentUser() user: CurrentUser,
+    @Args('workspaceId') workspaceId: string,
+    @Args('sessionId') sessionId: string
+  ) {
+    const lockFlag = `${COPILOT_LOCKER}:context:${sessionId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      return new TooManyRequest('Server is busy');
+    }
+    await this.checkChatSession(user, sessionId, workspaceId);
+
+    const context = await this.context.create(sessionId);
+    return context.id;
+  }
+}
+
+@Throttle()
+@Resolver(() => CopilotContextType)
+export class CopilotContextResolver {
+  constructor(
+    private readonly mutex: RequestMutex,
+
+    private readonly context: CopilotContextService
+  ) {}
+
+  @ResolveField(() => [CopilotContextDoc], {
+    description: 'list files in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async docs(@Parent() context: CopilotContextType): Promise<ContextDoc[]> {
+    const session = await this.context.get(context.id);
+    return session.listDocs();
+  }
+
+  @Mutation(() => [CopilotContextListItem], {
+    description: 'add a doc to context',
+  })
+  @CallMetric('ai', 'context_doc_add')
+  async addContextDoc(
+    @Args({ name: 'options', type: () => AddContextDocInput })
+    options: AddContextDocInput
+  ) {
+    const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      return new TooManyRequest('Server is busy');
+    }
+    const session = await this.context.get(options.contextId);
+
+    try {
+      return await session.addDocRecord(options.docId);
+    } catch (e: any) {
+      throw new CopilotFailedToModifyContext({
+        contextId: options.contextId,
+        message: e.message,
+      });
+    }
+  }
+
+  @Mutation(() => Boolean, {
+    description: 'remove a doc from context',
+  })
+  @CallMetric('ai', 'context_doc_remove')
+  async removeContextDoc(
+    @Args({ name: 'options', type: () => RemoveContextFileInput })
+    options: RemoveContextFileInput
+  ) {
+    const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      return new TooManyRequest('Server is busy');
+    }
+    const session = await this.context.get(options.contextId);
+
+    try {
+      return await session.removeDocRecord(options.fileId);
+    } catch (e: any) {
+      throw new CopilotFailedToModifyContext({
+        contextId: options.contextId,
+        message: e.message,
+      });
+    }
+  }
+
+  @ResolveField(() => [CopilotContextFile], {
+    description: 'list files in context',
+  })
+  @CallMetric('ai', 'context_file_list')
+  async files(
+    @Parent() context: CopilotContextType
+  ): Promise<CopilotContextFile[]> {
+    const session = await this.context.get(context.id);
+    return session.listFiles();
+  }
+}

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

@@ -0,0 +1,113 @@
+import { Injectable } from '@nestjs/common';
+import { PrismaClient } from '@prisma/client';
+
+import {
+  Cache,
+  CopilotInvalidContext,
+  CopilotSessionNotFound,
+} from '../../../base';
+import { ContextSession } from './session';
+import { ContextConfig, ContextConfigSchema } from './types';
+
+const CONTEXT_SESSION_KEY = 'context-session';
+
+@Injectable()
+export class CopilotContextService {
+  constructor(
+    private readonly cache: Cache,
+    private readonly db: PrismaClient
+  ) {}
+
+  private async saveConfig(
+    contextId: string,
+    config: ContextConfig,
+    refreshCache = false
+  ): Promise<void> {
+    if (!refreshCache) {
+      await this.db.aiContext.update({
+        where: { id: contextId },
+        data: { config },
+      });
+    }
+    await this.cache.set(`${CONTEXT_SESSION_KEY}:${contextId}`, config);
+  }
+
+  private async getCachedSession(
+    contextId: string
+  ): Promise<ContextSession | undefined> {
+    const cachedSession = await this.cache.get(
+      `${CONTEXT_SESSION_KEY}:${contextId}`
+    );
+    if (cachedSession) {
+      const config = ContextConfigSchema.safeParse(cachedSession);
+      if (config.success) {
+        return new ContextSession(
+          contextId,
+          config.data,
+          this.saveConfig.bind(this, contextId)
+        );
+      }
+    }
+    return undefined;
+  }
+
+  // NOTE: we only cache config to avoid frequent database queries
+  // but we do not need to cache session instances because a distributed
+  // lock is already apply to mutation operation for the same context in
+  // the resolver, so there will be no simultaneous writing to the config
+  private async cacheSession(
+    contextId: string,
+    config: ContextConfig
+  ): Promise<ContextSession> {
+    const dispatcher = this.saveConfig.bind(this, contextId);
+    await dispatcher(config, true);
+    return new ContextSession(contextId, config, dispatcher);
+  }
+
+  async create(sessionId: string): Promise<ContextSession> {
+    const session = await this.db.aiSession.findFirst({
+      where: { id: sessionId },
+      select: { workspaceId: true },
+    });
+    if (!session) {
+      throw new CopilotSessionNotFound();
+    }
+
+    // keep the context unique per session
+    const existsContext = await this.getBySessionId(sessionId);
+    if (existsContext) return existsContext;
+
+    const context = await this.db.aiContext.create({
+      data: {
+        sessionId,
+        config: { workspaceId: session.workspaceId, docs: [], files: [] },
+      },
+    });
+
+    const config = ContextConfigSchema.parse(context.config);
+    return await this.cacheSession(context.id, config);
+  }
+
+  async get(id: string): Promise<ContextSession> {
+    const context = await this.getCachedSession(id);
+    if (context) return context;
+    const ret = await this.db.aiContext.findUnique({
+      where: { id },
+      select: { config: true },
+    });
+    if (ret) {
+      const config = ContextConfigSchema.safeParse(ret.config);
+      if (config.success) return this.cacheSession(id, config.data);
+    }
+    throw new CopilotInvalidContext({ contextId: id });
+  }
+
+  async getBySessionId(sessionId: string): Promise<ContextSession | null> {
+    const existsContext = await this.db.aiContext.findFirst({
+      where: { sessionId },
+      select: { id: true },
+    });
+    if (existsContext) return this.get(existsContext.id);
+    return null;
+  }
+}

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

@@ -0,0 +1,58 @@
+import { ContextConfig, ContextDoc, ContextList } from './types';
+
+export class ContextSession implements AsyncDisposable {
+  constructor(
+    private readonly contextId: string,
+    private readonly config: ContextConfig,
+    private readonly dispatcher?: (config: ContextConfig) => Promise<void>
+  ) {}
+
+  get id() {
+    return this.contextId;
+  }
+
+  get workspaceId() {
+    return this.config.workspaceId;
+  }
+
+  listDocs(): ContextDoc[] {
+    return [...this.config.docs];
+  }
+
+  listFiles() {
+    return this.config.files.map(f => ({ ...f }));
+  }
+
+  get sortedList(): ContextList {
+    const { docs, files } = this.config;
+    return [...docs, ...files].toSorted(
+      (a, b) => a.createdAt - b.createdAt
+    ) as ContextList;
+  }
+
+  async addDocRecord(docId: string): Promise<ContextList> {
+    if (!this.config.docs.some(f => f.id === docId)) {
+      this.config.docs.push({ id: docId, createdAt: Date.now() });
+      await this.save();
+    }
+    return this.sortedList;
+  }
+
+  async removeDocRecord(docId: string): Promise<boolean> {
+    const index = this.config.docs.findIndex(f => f.id === docId);
+    if (index >= 0) {
+      this.config.docs.splice(index, 1);
+      await this.save();
+      return true;
+    }
+    return false;
+  }
+
+  async save() {
+    await this.dispatcher?.(this.config);
+  }
+
+  async [Symbol.asyncDispose]() {
+    await this.save();
+  }
+}

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

@@ -0,0 +1,69 @@
+import { z } from 'zod';
+
+declare global {
+  interface Events {
+    'workspace.doc.embedding': {
+      workspaceId: string;
+      docId: string;
+    };
+  }
+}
+
+export enum ContextFileStatus {
+  processing = 'processing',
+  finished = 'finished',
+  failed = 'failed',
+}
+
+export const ContextConfigSchema = z.object({
+  workspaceId: z.string(),
+  files: z
+    .object({
+      id: z.string(),
+      chunkSize: z.number(),
+      name: z.string(),
+      status: z.enum([
+        ContextFileStatus.processing,
+        ContextFileStatus.finished,
+        ContextFileStatus.failed,
+      ]),
+      blobId: z.string(),
+      createdAt: z.number(),
+    })
+    .array(),
+  docs: z
+    .object({
+      id: z.string(),
+      createdAt: z.number(),
+    })
+    .array(),
+});
+
+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[];
+
+export type ChunkSimilarity = {
+  chunk: number;
+  content: string;
+  distance: number | null;
+};
+
+export type FileChunkSimilarity = ChunkSimilarity & {
+  fileId: string;
+};
+
+export type DocChunkSimilarity = ChunkSimilarity & {
+  docId: string;
+};
+
+export type Embedding = {
+  /**
+   * The index of the embedding in the list of embeddings.
+   */
+  index: number;
+  content: string;
+  embedding: Array<number>;
+};

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

@@ -0,0 +1,11 @@
+export class GqlSignal implements AsyncDisposable {
+  readonly abortController = new AbortController();
+
+  get signal() {
+    return this.abortController.signal;
+  }
+
+  async [Symbol.asyncDispose]() {
+    this.abortController.abort();
+  }
+}

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

@@ -5,6 +5,11 @@ import { FeatureModule } from '../../core/features';
 import { PermissionModule } from '../../core/permission';
 import { QuotaModule } from '../../core/quota';
 import { Plugin } from '../registry';
+import {
+  CopilotContextResolver,
+  CopilotContextRootResolver,
+  CopilotContextService,
+} from './context';
 import { CopilotController } from './controller';
 import { ChatMessageCache } from './message';
 import { PromptService } from './prompt';
@@ -41,8 +46,13 @@ registerCopilotProvider(PerplexityProvider);
     CopilotProviderService,
     CopilotStorage,
     PromptsManagementResolver,
+    // workflow
     CopilotWorkflowService,
     ...CopilotWorkflowExecutors,
+    // context
+    CopilotContextRootResolver,
+    CopilotContextResolver,
+    CopilotContextService,
   ],
   controllers: [CopilotController],
   contributesTo: ServerFeature.Copilot,

packages/backend/server/src/plugins/copilot/providers/openai.ts

@@ -59,6 +59,7 @@ export class OpenAIProvider
 
   private readonly logger = new Logger(OpenAIProvider.type);
   private readonly instance: OpenAI;
+
   private existsModels: string[] | undefined;
 
   constructor(config: ClientOptions) {

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

@@ -23,7 +23,7 @@ import {
   CallMetric,
   CopilotFailedToCreateMessage,
   CopilotSessionNotFound,
-  FileUpload,
+  type FileUpload,
   RequestMutex,
   Throttle,
   TooManyRequest,

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

@@ -198,6 +198,13 @@ const CopilotImageOptionsSchema = CopilotProviderOptionsSchema.merge(
 
 export type CopilotImageOptions = z.infer<typeof CopilotImageOptionsSchema>;
 
+export type CopilotContextFile = {
+  id: string; // fileId
+  created_at: number;
+  // embedding status
+  status: 'in_progress' | 'completed' | 'failed';
+};
+
 export interface CopilotProvider {
   readonly type: CopilotProviderType;
   getCapabilities(): CopilotCapability[];

packages/backend/server/src/schema.gql

@@ -2,6 +2,11 @@
 # THIS FILE WAS AUTOMATICALLY GENERATED (DO NOT MODIFY)
 # ------------------------------------------------------
 
+input AddContextDocInput {
+  contextId: String!
+  docId: String!
+}
+
 type AlreadyInSpaceDataType {
   spaceId: String!
 }
@@ -25,12 +30,21 @@ type ChatMessage {
   role: String!
 }
 
+enum ContextFileStatus {
+  failed
+  finished
+  processing
+}
+
 type Copilot {
   """Get the session list of actions in the workspace"""
   actions: [String!]!
 
   """Get the session list of chats in the workspace"""
   chats: [String!]!
+
+  """Get the context list of a session"""
+  contexts(contextId: String, sessionId: String!): [CopilotContext!]!
   histories(docId: String, options: QueryChatHistoriesInput): [CopilotHistories!]!
 
   """Get the quota of the user in the workspace"""
@@ -38,6 +52,55 @@ type Copilot {
   workspaceId: ID
 }
 
+type CopilotContext {
+  """list files in context"""
+  docs: [CopilotContextDoc!]!
+
+  """list files in context"""
+  files: [CopilotContextFile!]!
+  id: ID!
+  workspaceId: String!
+}
+
+type CopilotContextDoc {
+  createdAt: SafeInt!
+  id: ID!
+}
+
+type CopilotContextFile {
+  blobId: String!
+  chunkSize: SafeInt!
+  createdAt: SafeInt!
+  id: ID!
+  name: String!
+  status: ContextFileStatus!
+}
+
+type CopilotContextFileNotSupportedDataType {
+  fileName: String!
+  message: String!
+}
+
+type CopilotContextListItem {
+  blobId: String
+  chunkSize: SafeInt
+  createdAt: SafeInt!
+  id: ID!
+  name: String
+  status: ContextFileStatus
+}
+
+type CopilotFailedToMatchContextDataType {
+  content: String!
+  contextId: String!
+  message: String!
+}
+
+type CopilotFailedToModifyContextDataType {
+  contextId: String!
+  message: String!
+}
+
 type CopilotHistories {
   """An mark identifying which view to use to display the session"""
   action: String
@@ -49,6 +112,10 @@ type CopilotHistories {
   tokens: Int!
 }
 
+type CopilotInvalidContextDataType {
+  contextId: String!
+}
+
 type CopilotMessageNotFoundDataType {
   messageId: String!
 }
@@ -244,7 +311,7 @@ type EditorType {
   name: String!
 }
 
-union ErrorDataUnion = AlreadyInSpaceDataType | BlobNotFoundDataType | CopilotMessageNotFoundDataType | CopilotPromptNotFoundDataType | CopilotProviderSideErrorDataType | DocAccessDeniedDataType | DocHistoryNotFoundDataType | DocNotFoundDataType | ExpectToGrantDocUserRolesDataType | ExpectToRevokeDocUserRolesDataType | ExpectToUpdateDocUserRoleDataType | InvalidEmailDataType | InvalidHistoryTimestampDataType | InvalidLicenseUpdateParamsDataType | InvalidPasswordLengthDataType | InvalidRuntimeConfigTypeDataType | MemberNotFoundInSpaceDataType | MissingOauthQueryParameterDataType | NotInSpaceDataType | QueryTooLongDataType | RuntimeConfigNotFoundDataType | SameSubscriptionRecurringDataType | SpaceAccessDeniedDataType | SpaceNotFoundDataType | SpaceOwnerNotFoundDataType | SpaceShouldHaveOnlyOneOwnerDataType | SubscriptionAlreadyExistsDataType | SubscriptionNotExistsDataType | SubscriptionPlanNotFoundDataType | UnknownOauthProviderDataType | UnsupportedSubscriptionPlanDataType | VersionRejectedDataType | WorkspaceMembersExceedLimitToDowngradeDataType | WorkspacePermissionNotFoundDataType | WrongSignInCredentialsDataType
+union ErrorDataUnion = AlreadyInSpaceDataType | BlobNotFoundDataType | CopilotContextFileNotSupportedDataType | CopilotFailedToMatchContextDataType | CopilotFailedToModifyContextDataType | CopilotInvalidContextDataType | CopilotMessageNotFoundDataType | CopilotPromptNotFoundDataType | CopilotProviderSideErrorDataType | DocAccessDeniedDataType | DocHistoryNotFoundDataType | DocNotFoundDataType | ExpectToGrantDocUserRolesDataType | ExpectToRevokeDocUserRolesDataType | ExpectToUpdateDocUserRoleDataType | InvalidEmailDataType | InvalidHistoryTimestampDataType | InvalidLicenseUpdateParamsDataType | InvalidPasswordLengthDataType | InvalidRuntimeConfigTypeDataType | MemberNotFoundInSpaceDataType | MissingOauthQueryParameterDataType | NotInSpaceDataType | QueryTooLongDataType | RuntimeConfigNotFoundDataType | SameSubscriptionRecurringDataType | SpaceAccessDeniedDataType | SpaceNotFoundDataType | SpaceOwnerNotFoundDataType | SpaceShouldHaveOnlyOneOwnerDataType | SubscriptionAlreadyExistsDataType | SubscriptionNotExistsDataType | SubscriptionPlanNotFoundDataType | UnknownOauthProviderDataType | UnsupportedSubscriptionPlanDataType | VersionRejectedDataType | WorkspaceMembersExceedLimitToDowngradeDataType | WorkspacePermissionNotFoundDataType | WrongSignInCredentialsDataType
 
 enum ErrorNames {
   ACCESS_DENIED
@@ -260,8 +327,12 @@ enum ErrorNames {
   CAN_NOT_BATCH_GRANT_DOC_OWNER_PERMISSIONS
   CAPTCHA_VERIFICATION_FAILED
   COPILOT_ACTION_TAKEN
+  COPILOT_CONTEXT_FILE_NOT_SUPPORTED
   COPILOT_FAILED_TO_CREATE_MESSAGE
   COPILOT_FAILED_TO_GENERATE_TEXT
+  COPILOT_FAILED_TO_MATCH_CONTEXT
+  COPILOT_FAILED_TO_MODIFY_CONTEXT
+  COPILOT_INVALID_CONTEXT
   COPILOT_MESSAGE_NOT_FOUND
   COPILOT_PROMPT_INVALID
   COPILOT_PROMPT_NOT_FOUND
@@ -574,6 +645,9 @@ type MissingOauthQueryParameterDataType {
 type Mutation {
   acceptInviteById(inviteId: String!, sendAcceptMail: Boolean, workspaceId: String!): Boolean!
   activateLicense(license: String!, workspaceId: String!): License!
+
+  """add a doc to context"""
+  addContextDoc(options: AddContextDocInput!): [CopilotContextListItem!]!
   addWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Boolean!
   approveMember(userId: String!, workspaceId: String!): String!
   cancelSubscription(idempotencyKey: String @deprecated(reason: "use header `Idempotency-Key`"), plan: SubscriptionPlan = Pro, workspaceId: String): SubscriptionType!
@@ -589,6 +663,9 @@ type Mutation {
   """Create a subscription checkout link of stripe"""
   createCheckoutSession(input: CreateCheckoutSessionInput!): String!
 
+  """Create a context session"""
+  createCopilotContext(sessionId: String!, workspaceId: String!): String!
+
   """Create a chat message"""
   createCopilotMessage(options: CreateChatMessageInput!): String!
 
@@ -631,6 +708,9 @@ type Mutation {
 
   """Remove user avatar"""
   removeAvatar: RemoveAvatar!
+
+  """remove a doc from context"""
+  removeContextDoc(options: RemoveContextFileInput!): Boolean!
   removeWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Boolean!
   resumeSubscription(idempotencyKey: String @deprecated(reason: "use header `Idempotency-Key`"), plan: SubscriptionPlan = Pro, workspaceId: String): SubscriptionType!
   revoke(userId: String!, workspaceId: String!): Boolean!
@@ -808,6 +888,11 @@ type RemoveAvatar {
   success: Boolean!
 }
 
+input RemoveContextFileInput {
+  contextId: String!
+  fileId: String!
+}
+
 input RevokeDocUserRoleInput {
   docId: String!
   userId: String!