From 631c8b8145d30789db05d595238070face8ae901 Mon Sep 17 00:00:00 2001
From: akumatus <akumatus@gmail.com>
Date: Fri, 14 Feb 2025 14:48:46 +0000
Subject: [PATCH] refactor(core): lazy create copilot session and context
 (#10170)

---
 .../ai/_common/components/chat-action-list.ts |   7 +-
 .../ai/_common/components/copy-more.ts        |   7 +-
 .../presets/ai/chat-panel/chat-panel-chips.ts |  16 +-
 .../presets/ai/chat-panel/chat-panel-input.ts |  89 +++++-----
 .../ai/chat-panel/chat-panel-messages.ts      |  11 +-
 .../ai/chat-panel/components/doc-chip.ts      |   4 +-
 .../blocksuite/presets/ai/chat-panel/index.ts | 156 +++++++++++-------
 .../ai/peek-view/chat-block-peek-view.ts      |   8 +-
 .../block-suite-editor/ai/request.ts          |  31 +---
 .../block-suite-editor/ai/setup-provider.tsx  |  15 +-
 .../affine-cloud-copilot/e2e/copilot.spec.ts  |   5 +-
 11 files changed, 179 insertions(+), 170 deletions(-)

diff --git a/packages/frontend/core/src/blocksuite/presets/ai/_common/components/chat-action-list.ts b/packages/frontend/core/src/blocksuite/presets/ai/_common/components/chat-action-list.ts
index 70bde30114..fe713c8d6e 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/_common/components/chat-action-list.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/_common/components/chat-action-list.ts
@@ -83,7 +83,7 @@ export class ChatActionList extends LitElement {
   accessor content: string = '';
 
   @property({ attribute: false })
-  accessor chatSessionId: string | undefined = undefined;
+  accessor getSessionId!: () => Promise<string | undefined>;
 
   @property({ attribute: false })
   accessor messageId: string | undefined = undefined;
@@ -100,7 +100,7 @@ export class ChatActionList extends LitElement {
       return nothing;
     }
 
-    const { host, content, chatSessionId, messageId, layoutDirection } = this;
+    const { host, content, messageId, layoutDirection } = this;
     const classes = classMap({
       'actions-container': true,
       horizontal: layoutDirection === 'horizontal',
@@ -138,11 +138,12 @@ export class ChatActionList extends LitElement {
                     blocks: this._currentBlockSelections,
                     images: this._currentImageSelections,
                   };
+                  const sessionId = await this.getSessionId();
                   const success = await action.handler(
                     host,
                     content,
                     currentSelections,
-                    chatSessionId,
+                    sessionId,
                     messageId
                   );
                   if (success) {
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/_common/components/copy-more.ts b/packages/frontend/core/src/blocksuite/presets/ai/_common/components/copy-more.ts
index 82fea70677..851054289d 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/_common/components/copy-more.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/_common/components/copy-more.ts
@@ -115,7 +115,7 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
   accessor content!: string;
 
   @property({ attribute: false })
-  accessor chatSessionId: string | undefined = undefined;
+  accessor getSessionId!: () => Promise<string | undefined>;
 
   @property({ attribute: false })
   accessor messageId: string | undefined = undefined;
@@ -162,7 +162,7 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
   }
 
   override render() {
-    const { host, content, isLast, messageId, chatSessionId, actions } = this;
+    const { host, content, isLast, messageId, actions } = this;
     return html`<style>
         .copy-more {
           margin-top: ${this.withMargin ? '8px' : '0px'};
@@ -217,11 +217,12 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
                 };
                 return html`<div
                   @click=${async () => {
+                    const sessionId = await this.getSessionId();
                     const success = await action.handler(
                       host,
                       content,
                       currentSelections,
-                      chatSessionId,
+                      sessionId,
                       messageId
                     );
 
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-chips.ts b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-chips.ts
index f2773f33f5..9789ffe6fc 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-chips.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-chips.ts
@@ -47,7 +47,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
   accessor chatContextValue!: ChatContextValue;
 
   @property({ attribute: false })
-  accessor chatContextId!: string | undefined;
+  accessor getContextId!: () => Promise<string | undefined>;
 
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<ChatContextValue>) => void;
@@ -196,34 +196,36 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
   };
 
   private readonly _addToContext = async (chip: ChatChip) => {
-    if (!AIProvider.context || !this.chatContextId) {
+    const contextId = await this.getContextId();
+    if (!contextId || !AIProvider.context) {
       return;
     }
     if (isDocChip(chip)) {
       await AIProvider.context.addContextDoc({
-        contextId: this.chatContextId,
+        contextId,
         docId: chip.docId,
       });
     } else {
       await AIProvider.context.addContextFile({
-        contextId: this.chatContextId,
+        contextId,
         fileId: chip.fileId,
       });
     }
   };
 
   private readonly _removeFromContext = async (chip: ChatChip) => {
-    if (!AIProvider.context || !this.chatContextId) {
+    const contextId = await this.getContextId();
+    if (!contextId || !AIProvider.context) {
       return;
     }
     if (isDocChip(chip)) {
       await AIProvider.context.removeContextDoc({
-        contextId: this.chatContextId,
+        contextId,
         docId: chip.docId,
       });
     } else {
       await AIProvider.context.removeContextFile({
-        contextId: this.chatContextId,
+        contextId,
         fileId: chip.fileId,
       });
     }
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts
index 9866b0cfb0..6ad48b6699 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts
@@ -260,7 +260,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
   accessor chatContextValue!: ChatContextValue;
 
   @property({ attribute: false })
-  accessor chatSessionId!: string | undefined;
+  accessor getSessionId!: () => Promise<string | undefined>;
 
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<ChatContextValue>) => void;
@@ -300,11 +300,9 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
   private async _updatePromptName() {
     if (this._lastPromptName !== this._promptName) {
       this._lastPromptName = this._promptName;
-      if (this.chatSessionId) {
-        await AIProvider.session?.updateSession(
-          this.chatSessionId,
-          this._promptName
-        );
+      const sessionId = await this.getSessionId();
+      if (sessionId) {
+        await AIProvider.session?.updateSession(sessionId, this._promptName);
       }
     }
   }
@@ -559,45 +557,46 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
     if (status === 'loading' || status === 'transmitting') return;
     if (!text) return;
 
-    const { images } = this.chatContextValue;
-    const { doc } = this.host;
-
-    this.updateContext({
-      images: [],
-      status: 'loading',
-      error: null,
-      quote: '',
-      markdown: '',
-    });
-
-    await this._updatePromptName();
-
-    const attachments = await Promise.all(
-      images?.map(image => readBlobAsURL(image))
-    );
-
-    const userInput = (markdown ? `${markdown}\n` : '') + text;
-    this.updateContext({
-      items: [
-        ...this.chatContextValue.items,
-        {
-          id: '',
-          role: 'user',
-          content: userInput,
-          createdAt: new Date().toISOString(),
-          attachments,
-        },
-        {
-          id: '',
-          role: 'assistant',
-          content: '',
-          createdAt: new Date().toISOString(),
-        },
-      ],
-    });
-
     try {
+      const { images } = this.chatContextValue;
+      const { doc } = this.host;
+
+      this.updateContext({
+        images: [],
+        status: 'loading',
+        error: null,
+        quote: '',
+        markdown: '',
+      });
+
+      const attachments = await Promise.all(
+        images?.map(image => readBlobAsURL(image))
+      );
+
+      const userInput = (markdown ? `${markdown}\n` : '') + text;
+      this.updateContext({
+        items: [
+          ...this.chatContextValue.items,
+          {
+            id: '',
+            role: 'user',
+            content: userInput,
+            createdAt: new Date().toISOString(),
+            attachments,
+          },
+          {
+            id: '',
+            role: 'assistant',
+            content: '',
+            createdAt: new Date().toISOString(),
+          },
+        ],
+      });
+
+      await this._updatePromptName();
+
       const abortController = new AbortController();
+      const sessionId = await this.getSessionId();
       const docs: DocContext[] = chips
         .filter(isDocChip)
         .filter(chip => !!chip.markdown?.value && chip.state === 'success')
@@ -606,7 +605,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
           markdown: chip.markdown?.value || '',
         }));
       const stream = AIProvider.actions.chat?.({
-        sessionId: this.chatSessionId,
+        sessionId,
         input: userInput,
         docs: docs,
         docId: doc.id,
@@ -638,7 +637,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
           const historyIds = await AIProvider.histories?.ids(
             doc.workspace.id,
             doc.id,
-            { sessionId: this.chatSessionId }
+            { sessionId }
           );
           if (!historyIds || !historyIds[0]) return;
           last.id = historyIds[0].messages.at(-1)?.id ?? '';
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-messages.ts b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-messages.ts
index d1c954a94f..67e95fb72f 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-messages.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-messages.ts
@@ -134,7 +134,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
   accessor chatContextValue!: ChatContextValue;
 
   @property({ attribute: false })
-  accessor chatSessionId!: string | undefined;
+  accessor getSessionId!: () => Promise<string | undefined>;
 
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<ChatContextValue>) => void;
@@ -415,7 +415,8 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
   retry = async () => {
     const { doc } = this.host;
     try {
-      if (!this.chatSessionId) return;
+      const sessionId = await this.getSessionId();
+      if (!sessionId) return;
 
       const abortController = new AbortController();
       const items = [...this.chatContextValue.items];
@@ -427,7 +428,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
       this.updateContext({ items, status: 'loading', error: null });
 
       const stream = AIProvider.actions.chat?.({
-        sessionId: this.chatSessionId,
+        sessionId,
         retry: true,
         docId: doc.id,
         workspaceId: doc.workspace.id,
@@ -482,7 +483,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
         .actions=${actions}
         .content=${content}
         .isLast=${isLast}
-        .chatSessionId=${this.chatSessionId}
+        .getSessionId=${this.getSessionId}
         .messageId=${messageId}
         .withMargin=${true}
         .retry=${() => this.retry()}
@@ -492,7 +493,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
             .actions=${actions}
             .host=${host}
             .content=${content}
-            .chatSessionId=${this.chatSessionId}
+            .getSessionId=${this.getSessionId}
             .messageId=${messageId ?? undefined}
             .withMargin=${true}
           ></chat-action-list>`
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts
index 91c1d22d71..c9f54befee 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts
@@ -63,8 +63,8 @@ export class ChatPanelDocChip extends SignalWatcher(
     super.updated(changedProperties);
     if (
       changedProperties.has('chip') &&
-      changedProperties.get('chip')?.state === 'candidate' &&
-      this.chip.state === 'processing'
+      this.chip.state === 'processing' &&
+      !this.chip.markdown
     ) {
       this.processDocChip().catch(console.error);
     }
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts
index e7d0602b87..8a02a2cc85 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts
@@ -36,6 +36,17 @@ import type {
 import type { ChatPanelMessages } from './chat-panel-messages';
 import { isDocContext } from './components/utils';
 
+const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
+  quote: '',
+  images: [],
+  abortController: null,
+  items: [],
+  chips: [],
+  status: 'idle',
+  error: null,
+  markdown: '',
+};
+
 export class ChatPanel extends WithDisposable(ShadowlessElement) {
   static override styles = css`
     chat-panel {
@@ -156,49 +167,80 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
   };
 
   private readonly _updateChips = async () => {
-    if (!this._chatSessionId || !this._chatContextId) return;
-
     const candidateChip: DocChip = {
       docId: this.doc.id,
       state: 'candidate',
     };
-    let chips: (DocChip | FileChip)[] = [];
-    if (this._chatContextId) {
-      const { docs = [], files = [] } =
-        (await AIProvider.context?.getContextDocsAndFiles(
-          this.doc.workspace.id,
-          this._chatSessionId,
-          this._chatContextId
-        )) || {};
-      const list = [...docs, ...files].sort(
-        (a, b) =>
-          new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
-      );
-      chips = list.map(item => {
-        let chip: DocChip | FileChip;
-        if (isDocContext(item)) {
-          chip = {
-            docId: item.id,
-            state: 'processing',
-          };
-        } else {
-          chip = {
-            fileId: item.id,
-            state: item.status === 'finished' ? 'success' : item.status,
-            fileName: item.name,
-            fileType: '',
-          };
-        }
-        return chip;
-      });
+
+    // context not initialized, show candidate chip
+    if (!this._chatSessionId || !this._chatContextId) {
+      this.chatContextValue = {
+        ...this.chatContextValue,
+        chips: [candidateChip],
+      };
+      return;
     }
 
+    // context initialized, show the chips
+    let chips: (DocChip | FileChip)[] = [];
+    const { docs = [], files = [] } =
+      (await AIProvider.context?.getContextDocsAndFiles(
+        this.doc.workspace.id,
+        this._chatSessionId,
+        this._chatContextId
+      )) || {};
+    const list = [...docs, ...files].sort(
+      (a, b) =>
+        new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
+    );
+    chips = list.map(item => {
+      let chip: DocChip | FileChip;
+      if (isDocContext(item)) {
+        chip = {
+          docId: item.id,
+          state: 'processing',
+        };
+      } else {
+        chip = {
+          fileId: item.id,
+          state: item.status === 'finished' ? 'success' : item.status,
+          fileName: item.name,
+          fileType: '',
+        };
+      }
+      return chip;
+    });
     this.chatContextValue = {
       ...this.chatContextValue,
       chips: chips.length === 0 ? [candidateChip] : chips,
     };
   };
 
+  private readonly _getSessionId = async () => {
+    if (this._chatSessionId) {
+      return this._chatSessionId;
+    }
+    this._chatSessionId = await AIProvider.session?.createSession(
+      this.doc.workspace.id,
+      this.doc.id
+    );
+    return this._chatSessionId;
+  };
+
+  private readonly _getContextId = async () => {
+    if (this._chatContextId) {
+      return this._chatContextId;
+    }
+    const sessionId = await this._getSessionId();
+    if (sessionId) {
+      this._chatContextId = await AIProvider.context?.createContext(
+        this.doc.workspace.id,
+        sessionId
+      );
+    }
+    return this._chatContextId;
+  };
+
   @property({ attribute: false })
   accessor host!: EditorHost;
 
@@ -221,20 +263,11 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
   accessor isLoading = false;
 
   @state()
-  accessor chatContextValue: ChatContextValue = {
-    quote: '',
-    images: [],
-    abortController: null,
-    items: [],
-    chips: [],
-    status: 'idle',
-    error: null,
-    markdown: '',
-  };
+  accessor chatContextValue: ChatContextValue = DEFAULT_CHAT_CONTEXT_VALUE;
 
-  private _chatSessionId: string | undefined;
+  private _chatSessionId: string | null | undefined = null;
 
-  private _chatContextId: string | undefined;
+  private _chatContextId: string | null | undefined = null;
 
   private readonly _scrollToEnd = () => {
     this._chatMessages.value?.scrollToEnd();
@@ -272,26 +305,29 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
     const userId = (await AIProvider.userInfo)?.id;
     if (!userId) return;
 
-    try {
-      this._chatSessionId = await AIProvider.session?.createSession(
-        this.doc.workspace.id,
-        this.doc.id
-      );
-      if (this._chatSessionId) {
-        this._chatContextId = await AIProvider.context?.createContext(
-          this.doc.workspace.id,
-          this._chatSessionId
-        );
-      }
-    } catch (e) {
-      console.error('init panel error', e);
+    const sessionIds = await AIProvider.session?.getSessionIds(
+      this.doc.workspace.id,
+      this.doc.id
+    );
+    if (sessionIds?.length) {
+      this._chatSessionId = sessionIds[0];
+      await this._updateHistory();
+    }
+    if (this._chatSessionId) {
+      this._chatContextId = await AIProvider.context?.getContextId(
+        this.doc.workspace.id,
+        this._chatSessionId
+      );
     }
-    await this._updateHistory();
     await this._updateChips();
   };
 
   protected override updated(_changedProperties: PropertyValues) {
     if (_changedProperties.has('doc')) {
+      this._chatSessionId = null;
+      this._chatContextId = null;
+      this.chatContextValue = DEFAULT_CHAT_CONTEXT_VALUE;
+
       requestAnimationFrame(async () => {
         await this._initPanel();
       });
@@ -376,7 +412,7 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
       <chat-panel-messages
         ${ref(this._chatMessages)}
         .chatContextValue=${this.chatContextValue}
-        .chatSessionId=${this._chatSessionId}
+        .getSessionId=${this._getSessionId}
         .updateContext=${this.updateContext}
         .host=${this.host}
         .isLoading=${this.isLoading}
@@ -385,14 +421,14 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
       <chat-panel-chips
         .host=${this.host}
         .chatContextValue=${this.chatContextValue}
-        .chatContextId=${this._chatContextId}
+        .getContextId=${this._getContextId}
         .updateContext=${this.updateContext}
         .docDisplayConfig=${this.docDisplayConfig}
         .docSearchMenuConfig=${this.docSearchMenuConfig}
       ></chat-panel-chips>
       <chat-panel-input
         .chatContextValue=${this.chatContextValue}
-        .chatSessionId=${this._chatSessionId}
+        .getSessionId=${this._getSessionId}
         .networkSearchConfig=${this.networkSearchConfig}
         .updateContext=${this.updateContext}
         .host=${this.host}
diff --git a/packages/frontend/core/src/blocksuite/presets/ai/peek-view/chat-block-peek-view.ts b/packages/frontend/core/src/blocksuite/presets/ai/peek-view/chat-block-peek-view.ts
index 47af871ea8..89b18a7fc1 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/peek-view/chat-block-peek-view.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/peek-view/chat-block-peek-view.ts
@@ -131,6 +131,10 @@ export class AIChatBlockPeekView extends LitElement {
     });
   };
 
+  private readonly _getSessionId = async () => {
+    return this.chatContext.currentSessionId ?? undefined;
+  };
+
   /**
    * Create a new AI chat block based on the current session and history messages
    */
@@ -408,7 +412,7 @@ export class AIChatBlockPeekView extends LitElement {
                 .actions=${actions}
                 .content=${message.content}
                 .isLast=${isLastReply}
-                .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
+                .getSessionId=${this._getSessionId}
                 .messageId=${message.id ?? undefined}
                 .retry=${() => this.retry()}
               ></chat-copy-more>`
@@ -418,7 +422,7 @@ export class AIChatBlockPeekView extends LitElement {
                 .host=${host}
                 .actions=${actions}
                 .content=${message.content}
-                .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
+                .getSessionId=${this._getSessionId}
                 .messageId=${message.id ?? undefined}
                 .layoutDirection=${'horizontal'}
               ></chat-action-list>`
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
index db1cc6828d..236b92be3b 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
@@ -30,30 +30,6 @@ export type ToImageOptions = TextToTextOptions & {
   seed?: string;
 };
 
-export async function createChatSession({
-  client,
-  workspaceId,
-  docId,
-  promptName = 'Chat With AFFiNE AI',
-}: {
-  client: CopilotClient;
-  workspaceId: string;
-  docId: string;
-  promptName?: string;
-}) {
-  const sessionId = await client.createSession({
-    workspaceId,
-    docId,
-    promptName,
-  });
-  // always update the prompt name
-  await client.updateSession({
-    sessionId,
-    promptName,
-  });
-  return sessionId;
-}
-
 async function resizeImage(blob: Blob | File): Promise<Blob | null> {
   let src = '';
   try {
@@ -90,7 +66,7 @@ async function createSessionMessage({
   client,
   docId,
   workspaceId,
-  promptName,
+  promptName = 'Chat With AFFiNE AI',
   content,
   sessionId: providedSessionId,
   attachments,
@@ -102,11 +78,10 @@ async function createSessionMessage({
   }
   const hasAttachments = attachments && attachments.length > 0;
   const sessionId = await (providedSessionId ??
-    createChatSession({
-      client,
+    client.createSession({
       workspaceId,
       docId,
-      promptName: promptName as string,
+      promptName,
     }));
 
   const options: Parameters<CopilotClient['createMessage']>[0] = {
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx
index 0208c64368..867516bed6 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx
@@ -11,7 +11,7 @@ import { z } from 'zod';
 
 import type { CopilotClient } from './copilot-client';
 import type { PromptKey } from './prompt';
-import { createChatSession, textToText, toImage } from './request';
+import { textToText, toImage } from './request';
 import { setupTracker } from './tracker';
 
 const filterStyleToPromptName = new Map(
@@ -37,13 +37,6 @@ export function setupAIProvider(
 ) {
   //#region actions
   AIProvider.provide('chat', options => {
-    const sessionId =
-      options.sessionId ??
-      createChatSession({
-        client,
-        workspaceId: options.workspaceId,
-        docId: options.docId,
-      });
     const { input, docs, ...rest } = options;
     const params = docs?.length
       ? {
@@ -58,7 +51,6 @@ export function setupAIProvider(
       ...rest,
       client,
       content: input,
-      sessionId,
       params,
     });
   });
@@ -408,10 +400,9 @@ Could you make a new website based on these notes and send back just the html fi
     createSession: async (
       workspaceId: string,
       docId: string,
-      promptName?: string
+      promptName = 'Chat With AFFiNE AI'
     ) => {
-      return createChatSession({
-        client,
+      return client.createSession({
         workspaceId,
         docId,
         promptName,
diff --git a/tests/affine-cloud-copilot/e2e/copilot.spec.ts b/tests/affine-cloud-copilot/e2e/copilot.spec.ts
index 4d233ec152..41973efc0e 100644
--- a/tests/affine-cloud-copilot/e2e/copilot.spec.ts
+++ b/tests/affine-cloud-copilot/e2e/copilot.spec.ts
@@ -104,12 +104,14 @@ const clearChat = async (page: Page) => {
 };
 
 const collectChat = async (page: Page) => {
+  await page.waitForTimeout(ONE_SECOND);
   const chatPanel = await page.waitForSelector('.chat-panel-messages');
   if (await chatPanel.$('.chat-panel-messages-placeholder')) {
     return [];
   }
   // wait ai response
   await page.waitForSelector('.chat-panel-messages .message chat-copy-more');
+  await page.waitForTimeout(ONE_SECOND);
   const lastMessage = await chatPanel.$$('.message').then(m => m[m.length - 1]);
   await lastMessage.waitForSelector('chat-copy-more');
   await page.waitForTimeout(ONE_SECOND);
@@ -489,7 +491,6 @@ test.describe('chat panel', () => {
     await page.getByTestId('chat-network-search').click();
     await typeChatSequentially(page, 'What is the weather in Shanghai today?');
     await page.keyboard.press('Enter');
-    await page.waitForTimeout(3000);
     let history = await collectChat(page);
     expect(history[0]).toEqual({
       name: 'You',
@@ -505,7 +506,6 @@ test.describe('chat panel', () => {
     await page.getByTestId('chat-network-search').click();
     await typeChatSequentially(page, 'What is the weather in Shanghai today?');
     await page.keyboard.press('Enter');
-    await page.waitForTimeout(3000);
     history = await collectChat(page);
     expect(history[0]).toEqual({
       name: 'You',
@@ -859,7 +859,6 @@ test.describe('chat with doc', () => {
 
     await typeChatSequentially(page, 'What is AFFiNE AI?');
     await page.keyboard.press('Enter');
-    await page.waitForTimeout(3000);
     const history = await collectChat(page);
     expect(history[0]).toEqual({
       name: 'You',