refactor(core): lazy create copilot session and context (#10170)

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) {

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
                     );
 

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,
       });
     }

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) {
+      const sessionId = await this.getSessionId();
-        await AIProvider.session?.updateSession(
+      if (sessionId) {
-          this.chatSessionId,
+        await AIProvider.session?.updateSession(sessionId, this._promptName);
-          this._promptName
-        );
       }
     }
   }
@@ -559,6 +557,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
     if (status === 'loading' || status === 'transmitting') return;
     if (!text) return;
 
+    try {
       const { images } = this.chatContextValue;
       const { doc } = this.host;
 
@@ -570,8 +569,6 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
         markdown: '',
       });
 
-    await this._updatePromptName();
-
       const attachments = await Promise.all(
         images?.map(image => readBlobAsURL(image))
       );
@@ -596,8 +593,10 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
         ],
       });
 
-    try {
+      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 ?? '';

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>`

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);
     }

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,14 +167,22 @@ 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',
     };
+
+    // 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)[] = [];
-    if (this._chatContextId) {
     const { docs = [], files = [] } =
       (await AIProvider.context?.getContextDocsAndFiles(
         this.doc.workspace.id,
@@ -191,14 +210,37 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
       }
       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 = {
+  accessor chatContextValue: ChatContextValue = DEFAULT_CHAT_CONTEXT_VALUE;
-    quote: '',
-    images: [],
-    abortController: null,
-    items: [],
-    chips: [],
-    status: 'idle',
-    error: null,
-    markdown: '',
-  };
 
-  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 {
+    const sessionIds = await AIProvider.session?.getSessionIds(
-      this._chatSessionId = await AIProvider.session?.createSession(
       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?.createContext(
+      this._chatContextId = await AIProvider.context?.getContextId(
         this.doc.workspace.id,
         this._chatSessionId
       );
     }
-    } catch (e) {
-      console.error('init panel error', e);
-    }
-    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}

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>`

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.createSession({
-      client,
       workspaceId,
       docId,
-      promptName: promptName as string,
+      promptName,
     }));
 
   const options: Parameters<CopilotClient['createMessage']>[0] = {

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({
+      return client.createSession({
-        client,
         workspaceId,
         docId,
         promptName,

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',