diff --git a/packages/backend/server/src/__tests__/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot.spec.ts
index 30ac64ede3..e3e7e57b9d 100644
--- a/packages/backend/server/src/__tests__/copilot.spec.ts
+++ b/packages/backend/server/src/__tests__/copilot.spec.ts
@@ -443,13 +443,13 @@ test('should be able to process message id', async t => {
   });
   const s = (await session.get(sessionId))!;
 
-  const textMessage = (await session.createMessage({
+  const textMessage = await session.createMessage({
     sessionId,
     content: 'hello',
-  }))!;
-  const anotherSessionMessage = (await session.createMessage({
+  });
+  const anotherSessionMessage = await session.createMessage({
     sessionId: 'another-session-id',
-  }))!;
+  });
 
   await t.notThrowsAsync(
     s.pushByMessageId(textMessage),
@@ -486,10 +486,10 @@ test('should be able to generate with message id', async t => {
     });
     const s = (await session.get(sessionId))!;
 
-    const message = (await session.createMessage({
+    const message = await session.createMessage({
       sessionId,
       content: 'hello',
-    }))!;
+    });
 
     await s.pushByMessageId(message);
     const finalMessages = s
@@ -508,10 +508,10 @@ test('should be able to generate with message id', async t => {
     });
     const s = (await session.get(sessionId))!;
 
-    const message = (await session.createMessage({
+    const message = await session.createMessage({
       sessionId,
       attachments: ['https://affine.pro/example.jpg'],
-    }))!;
+    });
 
     await s.pushByMessageId(message);
     const finalMessages = s
@@ -535,9 +535,9 @@ test('should be able to generate with message id', async t => {
     });
     const s = (await session.get(sessionId))!;
 
-    const message = (await session.createMessage({
+    const message = await session.createMessage({
       sessionId,
-    }))!;
+    });
 
     await s.pushByMessageId(message);
     const finalMessages = s
@@ -563,10 +563,10 @@ test('should save message correctly', async t => {
   });
   const s = (await session.get(sessionId))!;
 
-  const message = (await session.createMessage({
+  const message = await session.createMessage({
     sessionId,
     content: 'hello',
-  }))!;
+  });
 
   await s.pushByMessageId(message);
   t.is(s.stashMessages.length, 1, 'should get stash messages');
@@ -592,10 +592,10 @@ test('should revert message correctly', async t => {
     });
     const s = (await session.get(sessionId))!;
 
-    const message = (await session.createMessage({
+    const message = await session.createMessage({
       sessionId,
       content: '1',
-    }))!;
+    });
 
     await s.pushByMessageId(message);
     await s.save();
diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts
index 7117cb483e..ccb6fc33b2 100644
--- a/packages/backend/server/src/plugins/copilot/context/resolver.ts
+++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts
@@ -261,12 +261,12 @@ export class CopilotContextRootResolver {
     @CurrentUser() user: CurrentUser,
     @Args('sessionId', { nullable: true }) sessionId?: string,
     @Args('contextId', { nullable: true }) contextId?: string
-  ) {
+  ): Promise<CopilotContextType[]> {
     if (sessionId || contextId) {
       const lockFlag = `${COPILOT_LOCKER}:context:${sessionId || contextId}`;
       await using lock = await this.mutex.acquire(lockFlag);
       if (!lock) {
-        return new TooManyRequest('Server is busy');
+        throw new TooManyRequest('Server is busy');
       }
 
       if (contextId) {
@@ -294,11 +294,11 @@ export class CopilotContextRootResolver {
     @CurrentUser() user: CurrentUser,
     @Args('workspaceId') workspaceId: string,
     @Args('sessionId') sessionId: string
-  ) {
+  ): Promise<string> {
     const lockFlag = `${COPILOT_LOCKER}:context:${sessionId}`;
     await using lock = await this.mutex.acquire(lockFlag);
     if (!lock) {
-      return new TooManyRequest('Server is busy');
+      throw new TooManyRequest('Server is busy');
     }
     await this.checkChatSession(user, sessionId, workspaceId);
 
@@ -314,7 +314,7 @@ export class CopilotContextRootResolver {
     @CurrentUser() user: CurrentUser,
     @Args('workspaceId') workspaceId: string,
     @Args('docId', { type: () => [String] }) docIds: string[]
-  ) {
+  ): Promise<boolean> {
     await this.ac
       .user(user.id)
       .workspace(workspaceId)
@@ -339,7 +339,7 @@ export class CopilotContextRootResolver {
   async queryWorkspaceEmbeddingStatus(
     @CurrentUser() user: CurrentUser,
     @Args('workspaceId') workspaceId: string
-  ) {
+  ): Promise<ContextWorkspaceEmbeddingStatus> {
     await this.ac
       .user(user.id)
       .workspace(workspaceId)
@@ -386,7 +386,7 @@ export class CopilotContextResolver {
   @CallMetric('ai', 'context_file_list')
   async collections(
     @Parent() context: CopilotContextType
-  ): Promise<ContextCategory[]> {
+  ): Promise<CopilotContextCategory[]> {
     const session = await this.context.get(context.id);
     const collections = session.collections;
     await this.models.copilotContext.mergeDocStatus(
@@ -403,7 +403,7 @@ export class CopilotContextResolver {
   @CallMetric('ai', 'context_file_list')
   async tags(
     @Parent() context: CopilotContextType
-  ): Promise<ContextCategory[]> {
+  ): Promise<CopilotContextCategory[]> {
     const session = await this.context.get(context.id);
     const tags = session.tags;
     await this.models.copilotContext.mergeDocStatus(
@@ -485,11 +485,11 @@ export class CopilotContextResolver {
   async removeContextCategory(
     @Args({ name: 'options', type: () => RemoveContextCategoryInput })
     options: RemoveContextCategoryInput
-  ) {
+  ): Promise<boolean> {
     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);
 
@@ -545,11 +545,11 @@ export class CopilotContextResolver {
   async removeContextDoc(
     @Args({ name: 'options', type: () => RemoveContextDocInput })
     options: RemoveContextDocInput
-  ) {
+  ): Promise<boolean> {
     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);
 
@@ -574,7 +574,7 @@ export class CopilotContextResolver {
     options: AddContextFileInput,
     @Args({ name: 'content', type: () => GraphQLUpload })
     content: FileUpload
-  ) {
+  ): Promise<CopilotContextFile> {
     if (!this.context.canEmbedding) {
       throw new CopilotEmbeddingUnavailable();
     }
@@ -582,7 +582,7 @@ export class CopilotContextResolver {
     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 length = Number(ctx.req.headers['content-length']);
@@ -632,7 +632,7 @@ export class CopilotContextResolver {
   async removeContextFile(
     @Args({ name: 'options', type: () => RemoveContextFileInput })
     options: RemoveContextFileInput
-  ) {
+  ): Promise<boolean> {
     if (!this.context.canEmbedding) {
       throw new CopilotEmbeddingUnavailable();
     }
@@ -640,7 +640,7 @@ export class CopilotContextResolver {
     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);
 
diff --git a/packages/backend/server/src/plugins/copilot/message.ts b/packages/backend/server/src/plugins/copilot/message.ts
index 8fe953d839..4dd7e987ac 100644
--- a/packages/backend/server/src/plugins/copilot/message.ts
+++ b/packages/backend/server/src/plugins/copilot/message.ts
@@ -1,6 +1,6 @@
 import { randomUUID } from 'node:crypto';
 
-import { Injectable, Logger } from '@nestjs/common';
+import { Injectable } from '@nestjs/common';
 
 import { SessionCache } from '../../base';
 import { SubmittedMessage, SubmittedMessageSchema } from './types';
@@ -10,26 +10,18 @@ const CHAT_MESSAGE_TTL = 3600 * 1 * 1000; // 1 hours
 
 @Injectable()
 export class ChatMessageCache {
-  private readonly logger = new Logger(ChatMessageCache.name);
   constructor(private readonly cache: SessionCache) {}
 
   async get(id: string): Promise<SubmittedMessage | undefined> {
     return await this.cache.get(`${CHAT_MESSAGE_KEY}:${id}`);
   }
 
-  async set(message: SubmittedMessage): Promise<string | undefined> {
-    try {
-      const parsed = SubmittedMessageSchema.safeParse(message);
-      if (parsed.success) {
-        const id = randomUUID();
-        await this.cache.set(`${CHAT_MESSAGE_KEY}:${id}`, parsed.data, {
-          ttl: CHAT_MESSAGE_TTL,
-        });
-        return id;
-      }
-    } catch (e: any) {
-      this.logger.error(`Failed to get chat message from cache: ${e.message}`);
-    }
-    return undefined;
+  async set(message: SubmittedMessage): Promise<string> {
+    const parsedMessage = SubmittedMessageSchema.parse(message);
+    const id = randomUUID();
+    await this.cache.set(`${CHAT_MESSAGE_KEY}:${id}`, parsedMessage, {
+      ttl: CHAT_MESSAGE_TTL,
+    });
+    return id;
   }
 }
diff --git a/packages/backend/server/src/plugins/copilot/resolver.ts b/packages/backend/server/src/plugins/copilot/resolver.ts
index 7c9703b760..345aefaa65 100644
--- a/packages/backend/server/src/plugins/copilot/resolver.ts
+++ b/packages/backend/server/src/plugins/copilot/resolver.ts
@@ -169,7 +169,7 @@ class QueryChatHistoriesInput implements Partial<ListHistoriesOptions> {
 class ChatMessageType implements Partial<ChatMessage> {
   // id will be null if message is a prompt message
   @Field(() => ID, { nullable: true })
-  id!: string;
+  id!: string | undefined;
 
   @Field(() => String)
   role!: 'system' | 'assistant' | 'user';
@@ -310,7 +310,7 @@ export class CopilotResolver {
     description: 'Get the quota of the user in the workspace',
     complexity: 2,
   })
-  async getQuota(@CurrentUser() user: CurrentUser) {
+  async getQuota(@CurrentUser() user: CurrentUser): Promise<CopilotQuotaType> {
     return await this.chatSession.getQuota(user.id);
   }
 
@@ -324,8 +324,8 @@ export class CopilotResolver {
     @CurrentUser() user: CurrentUser,
     @Args('docId', { nullable: true }) docId?: string,
     @Args('options', { nullable: true }) options?: QueryChatSessionsInput
-  ) {
-    return await this.sessions(copilot, user, docId, options);
+  ): Promise<string[]> {
+    return (await this.sessions(copilot, user, docId, options)).map(s => s.id);
   }
 
   @ResolveField(() => [CopilotSessionType], {
@@ -337,7 +337,7 @@ export class CopilotResolver {
     @CurrentUser() user: CurrentUser,
     @Args('docId', { nullable: true }) docId?: string,
     @Args('options', { nullable: true }) options?: QueryChatSessionsInput
-  ) {
+  ): Promise<CopilotSessionType[]> {
     if (!copilot.workspaceId) return [];
     await this.ac
       .user(user.id)
@@ -359,7 +359,7 @@ export class CopilotResolver {
     @CurrentUser() user: CurrentUser,
     @Args('docId', { nullable: true }) docId?: string,
     @Args('options', { nullable: true }) options?: QueryChatHistoriesInput
-  ) {
+  ): Promise<CopilotHistoriesType[]> {
     const workspaceId = copilot.workspaceId;
     if (!workspaceId) {
       return [];
@@ -387,7 +387,9 @@ export class CopilotResolver {
     return histories.map(h => ({
       ...h,
       // filter out empty messages
-      messages: h.messages.filter(m => m.content || m.attachments?.length),
+      messages: h.messages.filter(
+        m => m.content || m.attachments?.length
+      ) as ChatMessageType[],
     }));
   }
 
@@ -399,12 +401,12 @@ export class CopilotResolver {
     @CurrentUser() user: CurrentUser,
     @Args({ name: 'options', type: () => CreateChatSessionInput })
     options: CreateChatSessionInput
-  ) {
+  ): Promise<string> {
     await this.ac.user(user.id).doc(options).allowLocal().assert('Doc.Update');
     const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
     await using lock = await this.mutex.acquire(lockFlag);
     if (!lock) {
-      return new TooManyRequest('Server is busy');
+      throw new TooManyRequest('Server is busy');
     }
 
     if (options.workspaceId === options.docId) {
@@ -428,7 +430,7 @@ export class CopilotResolver {
     @CurrentUser() user: CurrentUser,
     @Args({ name: 'options', type: () => UpdateChatSessionInput })
     options: UpdateChatSessionInput
-  ) {
+  ): Promise<string> {
     const session = await this.chatSession.get(options.sessionId);
     if (!session) {
       throw new CopilotSessionNotFound();
@@ -442,7 +444,7 @@ export class CopilotResolver {
     const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${workspaceId}`;
     await using lock = await this.mutex.acquire(lockFlag);
     if (!lock) {
-      return new TooManyRequest('Server is busy');
+      throw new TooManyRequest('Server is busy');
     }
 
     await this.chatSession.checkQuota(user.id);
@@ -460,12 +462,12 @@ export class CopilotResolver {
     @CurrentUser() user: CurrentUser,
     @Args({ name: 'options', type: () => ForkChatSessionInput })
     options: ForkChatSessionInput
-  ) {
+  ): Promise<string> {
     await this.ac.user(user.id).doc(options).allowLocal().assert('Doc.Update');
     const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
     await using lock = await this.mutex.acquire(lockFlag);
     if (!lock) {
-      return new TooManyRequest('Server is busy');
+      throw new TooManyRequest('Server is busy');
     }
 
     if (options.workspaceId === options.docId) {
@@ -489,15 +491,15 @@ export class CopilotResolver {
     @CurrentUser() user: CurrentUser,
     @Args({ name: 'options', type: () => DeleteSessionInput })
     options: DeleteSessionInput
-  ) {
+  ): Promise<string[]> {
     await this.ac.user(user.id).doc(options).allowLocal().assert('Doc.Update');
     if (!options.sessionIds.length) {
-      return new NotFoundException('Session not found');
+      throw new NotFoundException('Session not found');
     }
     const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${options.workspaceId}`;
     await using lock = await this.mutex.acquire(lockFlag);
     if (!lock) {
-      return new TooManyRequest('Server is busy');
+      throw new TooManyRequest('Server is busy');
     }
 
     return await this.chatSession.cleanup({
@@ -514,15 +516,15 @@ export class CopilotResolver {
     @CurrentUser() user: CurrentUser,
     @Args({ name: 'options', type: () => CreateChatMessageInput })
     options: CreateChatMessageInput
-  ) {
+  ): Promise<string> {
     const lockFlag = `${COPILOT_LOCKER}:message:${user?.id}:${options.sessionId}`;
     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.chatSession.get(options.sessionId);
     if (!session || session.config.userId !== user.id) {
-      return new BadRequestException('Session not found');
+      throw new BadRequestException('Session not found');
     }
 
     if (options.blobs) {
@@ -564,7 +566,7 @@ export class UserCopilotResolver {
   async copilot(
     @CurrentUser() user: CurrentUser,
     @Args('workspaceId', { nullable: true }) workspaceId?: string
-  ) {
+  ): Promise<CopilotType> {
     if (workspaceId) {
       await this.ac
         .user(user.id)
diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts
index e0a1f2007c..60ec974b03 100644
--- a/packages/backend/server/src/plugins/copilot/session.ts
+++ b/packages/backend/server/src/plugins/copilot/session.ts
@@ -713,7 +713,7 @@ export class ChatSessionService {
     });
   }
 
-  async createMessage(message: SubmittedMessage): Promise<string | undefined> {
+  async createMessage(message: SubmittedMessage): Promise<string> {
     return await this.messageCache.set(message);
   }