feat: add doc copilot context api (#10103)

### What Changed?
- Add graphql APIs.
- Provide context and session service in `AIProvider`.
- Rename the state from `embedding` to `processing`.
- Reafctor front-end session create, update and save logic.

Persist the document selected by the user:
[录屏2025-02-08 11.04.40.mov <span class="graphite__hidden">(uploaded via Graphite)</span> <img class="graphite__hidden" src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/195a85f2-43c4-4e49-88d9-6b5fc4f235ca.mov" />](https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/195a85f2-43c4-4e49-88d9-6b5fc4f235ca.mov)

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

@@ -36,12 +36,12 @@ class AddContextDocInput {
 }
 
 @InputType()
-class RemoveContextFileInput {
+class RemoveContextDocInput {
   @Field(() => String)
   contextId!: string;
 
   @Field(() => String)
-  fileId!: string;
+  docId!: string;
 }
 
 @ObjectType('CopilotContext')
@@ -227,8 +227,8 @@ export class CopilotContextResolver {
   })
   @CallMetric('ai', 'context_doc_remove')
   async removeContextDoc(
-    @Args({ name: 'options', type: () => RemoveContextFileInput })
+    @Args({ name: 'options', type: () => RemoveContextDocInput })
-    options: RemoveContextFileInput
+    options: RemoveContextDocInput
   ) {
     const lockFlag = `${COPILOT_LOCKER}:context:${options.contextId}`;
     await using lock = await this.mutex.acquire(lockFlag);
@@ -238,7 +238,7 @@ export class CopilotContextResolver {
     const session = await this.context.get(options.contextId);
 
     try {
-      return await session.removeDocRecord(options.fileId);
+      return await session.removeDocRecord(options.docId);
     } catch (e: any) {
       throw new CopilotFailedToModifyContext({
         contextId: options.contextId,

packages/backend/server/src/schema.gql

@@ -713,7 +713,7 @@ type Mutation {
   removeAvatar: RemoveAvatar!
 
   """remove a doc from context"""
-  removeContextDoc(options: RemoveContextFileInput!): Boolean!
+  removeContextDoc(options: RemoveContextDocInput!): 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!
@@ -891,9 +891,9 @@ type RemoveAvatar {
   success: Boolean!
 }
 
-input RemoveContextFileInput {
+input RemoveContextDocInput {
   contextId: String!
-  fileId: String!
+  docId: String!
 }
 
 input RevokeDocUserRoleInput {

packages/frontend/core/src/blocksuite/presets/ai/actions/doc-handler.ts

@@ -236,9 +236,14 @@ export function handleInlineAskAIAction(
     host.selection.set([selection]);
 
     selectAboveBlocks(host)
-      .then(context => {
+      .then(async context => {
-        assertExists(AIProvider.actions.chat);
+        if (!AIProvider.session || !AIProvider.actions.chat) return;
+        const sessionId = await AIProvider.session.createSession(
+          host.doc.workspace.id,
+          host.doc.id
+        );
         const stream = AIProvider.actions.chat({
+          sessionId,
           input: `${context}\n${input}`,
           stream: true,
           host,

packages/frontend/core/src/blocksuite/presets/ai/actions/types.ts

@@ -1,4 +1,9 @@
-import type { getCopilotHistoriesQuery, RequestOptions } from '@affine/graphql';
+import type {
+  CopilotContextDoc,
+  CopilotContextFile,
+  getCopilotHistoriesQuery,
+  RequestOptions,
+} from '@affine/graphql';
 import type { EditorHost } from '@blocksuite/affine/block-std';
 import type { GfxModel } from '@blocksuite/affine/block-std/gfx';
 import type { BlockModel } from '@blocksuite/affine/store';
@@ -105,10 +110,9 @@ declare global {
       T['stream'] extends true ? TextStream : Promise<string>;
 
     interface ChatOptions extends AITextActionOptions {
-      // related documents
-      docs?: DocContext[];
       sessionId?: string;
       isRootSession?: boolean;
+      docs?: DocContext[];
     }
 
     interface TranslateOptions extends AITextActionOptions {
@@ -232,6 +236,40 @@ declare global {
       ): AIActionTextResponse<T>;
     }
 
+    interface AIContextService {
+      createContext: (
+        workspaceId: string,
+        sessionId: string
+      ) => Promise<string>;
+      getContextId: (
+        workspaceId: string,
+        sessionId: string
+      ) => Promise<string | undefined>;
+      addContextDoc: (options: {
+        contextId: string;
+        docId: string;
+      }) => Promise<Array<{ id: string }>>;
+      removeContextDoc: (options: {
+        contextId: string;
+        docId: string;
+      }) => Promise<boolean>;
+      addContextFile: (options: {
+        contextId: string;
+        fileId: string;
+      }) => Promise<void>;
+      removeContextFile: (options: {
+        contextId: string;
+        fileId: string;
+      }) => Promise<void>;
+      getContextDocsAndFiles: (
+        workspaceId: string,
+        sessionId: string,
+        contextId: string
+      ) => Promise<
+        { docs: CopilotContextDoc[]; files: CopilotContextFile[] } | undefined
+      >;
+    }
+
     // TODO(@Peng): should be refactored to get rid of implement details (like messages, action, role, etc.)
     interface AIHistory {
       sessionId: string;
@@ -256,6 +294,15 @@ declare global {
       >[];
     };
 
+    interface AISessionService {
+      createSession: (
+        workspaceId: string,
+        docId: string,
+        promptName?: string
+      ) => Promise<string>;
+      updateSession: (sessionId: string, promptName: string) => Promise<string>;
+    }
+
     interface AIHistoryService {
       // non chat histories
       actions: (

packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-context.ts

@@ -46,7 +46,6 @@ export type ChatContextValue = {
   // chips of workspace doc or user uploaded file
   chips: ChatChip[];
   abortController: AbortController | null;
-  chatSessionId: string | null;
 };
 
 export type ChatBlockMessage = ChatMessage & {
@@ -55,20 +54,14 @@ export type ChatBlockMessage = ChatMessage & {
   avatarUrl?: string;
 };
 
-export type ChipState =
+export type ChipState = 'candidate' | 'processing' | 'success' | 'failed';
-  | 'candidate'
-  | 'uploading'
-  | 'embedding'
-  | 'success'
-  | 'failed';
 
 export interface BaseChip {
   /**
    * candidate: the chip is a candidate for the chat
-   * uploading: the chip is uploading
+   * processing: the chip is processing
-   * embedding: the chip is embedding
+   * success: the chip is successfully processed
-   * success: the chip is successfully embedded
+   * failed: the chip is failed to process
-   * failed: the chip is failed to embed
    */
   state: ChipState;
   tooltip?: string;

packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-chips.ts

@@ -10,6 +10,7 @@ import { css, html } from 'lit';
 import { property, query } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
 
+import { AIProvider } from '../provider';
 import type { DocDisplayConfig, DocSearchMenuConfig } from './chat-config';
 import type { BaseChip, ChatChip, ChatContextValue } from './chat-context';
 import { getChipKey, isDocChip, isFileChip } from './components/utils';
@@ -45,6 +46,9 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
   @property({ attribute: false })
   accessor chatContextValue!: ChatContextValue;
 
+  @property({ attribute: false })
+  accessor chatContextId!: string | undefined;
+
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<ChatContextValue>) => void;
 
@@ -69,6 +73,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
           if (isDocChip(chip)) {
             return html`<chat-panel-doc-chip
               .chip=${chip}
+              .addChip=${this._addChip}
               .updateChip=${this._updateChip}
               .removeChip=${this._removeChip}
               .docDisplayConfig=${this.docDisplayConfig}
@@ -120,11 +125,12 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
     });
   };
 
-  private readonly _addChip = (chip: ChatChip) => {
+  private readonly _addChip = async (chip: ChatChip) => {
     if (
       this.chatContextValue.chips.length === 1 &&
       this.chatContextValue.chips[0].state === 'candidate'
     ) {
+      await this._addToContext(chip);
       this.updateContext({
         chips: [chip],
       });
@@ -132,12 +138,16 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
     }
     // remove the chip if it already exists
     const chips = this.chatContextValue.chips.filter(item => {
-      if (isDocChip(item)) {
+      if (isDocChip(chip)) {
-        return !isDocChip(chip) || item.docId !== chip.docId;
+        return !isDocChip(item) || item.docId !== chip.docId;
       } else {
-        return !isFileChip(chip) || item.fileId !== chip.fileId;
+        return !isFileChip(item) || item.fileId !== chip.fileId;
       }
     });
+    if (chips.length < this.chatContextValue.chips.length) {
+      await this._removeFromContext(chip);
+    }
+    await this._addToContext(chip);
     this.updateContext({
       chips: [...chips, chip],
     });
@@ -167,15 +177,55 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
     });
   };
 
-  private readonly _removeChip = (chip: ChatChip) => {
+  private readonly _removeChip = async (chip: ChatChip) => {
-    this.updateContext({
+    if (isDocChip(chip)) {
-      chips: this.chatContextValue.chips.filter(item => {
+      await this._removeFromContext(chip);
-        if (isDocChip(item)) {
+      this.updateContext({
-          return !isDocChip(chip) || item.docId !== chip.docId;
+        chips: this.chatContextValue.chips.filter(item => {
-        } else {
+          return !isDocChip(item) || item.docId !== chip.docId;
-          return !isFileChip(chip) || item.fileId !== chip.fileId;
+        }),
-        }
+      });
-      }),
+    } else {
-    });
+      await this._removeFromContext(chip);
+      this.updateContext({
+        chips: this.chatContextValue.chips.filter(item => {
+          return !isFileChip(item) || item.fileId !== chip.fileId;
+        }),
+      });
+    }
+  };
+
+  private readonly _addToContext = async (chip: ChatChip) => {
+    if (!AIProvider.context || !this.chatContextId) {
+      return;
+    }
+    if (isDocChip(chip)) {
+      await AIProvider.context.addContextDoc({
+        contextId: this.chatContextId,
+        docId: chip.docId,
+      });
+    } else {
+      await AIProvider.context.addContextFile({
+        contextId: this.chatContextId,
+        fileId: chip.fileId,
+      });
+    }
+  };
+
+  private readonly _removeFromContext = async (chip: ChatChip) => {
+    if (!AIProvider.context || !this.chatContextId) {
+      return;
+    }
+    if (isDocChip(chip)) {
+      await AIProvider.context.removeContextDoc({
+        contextId: this.chatContextId,
+        docId: chip.docId,
+      });
+    } else {
+      await AIProvider.context.removeContextFile({
+        contextId: this.chatContextId,
+        fileId: chip.fileId,
+      });
+    }
   };
 }

packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts

@@ -258,6 +258,9 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
   @property({ attribute: false })
   accessor chatContextValue!: ChatContextValue;
 
+  @property({ attribute: false })
+  accessor chatSessionId!: string | undefined;
+
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<ChatContextValue>) => void;
 
@@ -267,6 +270,44 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
   @property({ attribute: false })
   accessor networkSearchConfig!: AINetworkSearchConfig;
 
+  private _lastPromptName: string | null = null;
+
+  private get _isNetworkActive() {
+    return (
+      !!this.networkSearchConfig.visible.value &&
+      !!this.networkSearchConfig.enabled.value
+    );
+  }
+
+  private get _isNetworkDisabled() {
+    return (
+      !!this.chatContextValue.images.length ||
+      !!this.chatContextValue.chips.filter(chip => chip.state !== 'candidate')
+        .length
+    );
+  }
+
+  private get _promptName() {
+    if (this._isNetworkDisabled) {
+      return 'Chat With AFFiNE AI';
+    }
+    return this._isNetworkActive
+      ? 'Search With AFFiNE AI'
+      : 'Chat With AFFiNE AI';
+  }
+
+  private async _updatePromptName() {
+    if (this._lastPromptName !== this._promptName) {
+      this._lastPromptName = this._promptName;
+      if (this.chatSessionId) {
+        await AIProvider.session?.updateSession(
+          this.chatSessionId,
+          this._promptName
+        );
+      }
+    }
+  }
+
   private _addImages(images: File[]) {
     const oldImages = this.chatContextValue.images;
     this.updateContext({
@@ -363,12 +404,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
     const { images, status } = this.chatContextValue;
     const hasImages = images.length > 0;
     const maxHeight = hasImages ? 272 + 2 : 200 + 2;
-    const networkDisabled =
+    const uploadDisabled = this._isNetworkActive && !this._isNetworkDisabled;
-      !!this.chatContextValue.images.length ||
-      !!this.chatContextValue.chips.filter(chip => chip.state !== 'candidate')
-        .length;
-    const networkActive = !!this.networkSearchConfig.enabled.value;
-    const uploadDisabled = networkActive && !networkDisabled;
     return html`<style>
         .chat-panel-input {
           border-color: ${this.focused
@@ -461,9 +497,9 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
                 <div
                   class="chat-network-search"
                   data-testid="chat-network-search"
-                  aria-disabled=${networkDisabled}
+                  aria-disabled=${this._isNetworkDisabled}
-                  data-active=${networkActive}
+                  data-active=${this._isNetworkActive}
-                  @click=${networkDisabled
+                  @click=${this._isNetworkDisabled
                     ? undefined
                     : this._toggleNetworkSearch}
                   @pointerdown=${stopPropagation}
@@ -520,11 +556,9 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
   send = async (text: string) => {
     const { status, markdown, chips } = this.chatContextValue;
     if (status === 'loading' || status === 'transmitting') return;
+    if (!text) return;
 
     const { images } = this.chatContextValue;
-    if (!text) {
-      return;
-    }
     const { doc } = this.host;
 
     this.updateContext({
@@ -535,6 +569,8 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
       markdown: '',
     });
 
+    await this._updatePromptName();
+
     const attachments = await Promise.all(
       images?.map(image => readBlobAsURL(image))
     );
@@ -569,6 +605,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
           markdown: chip.markdown?.value || '',
         }));
       const stream = AIProvider.actions.chat?.({
+        sessionId: this.chatSessionId,
         input: userInput,
         docs: docs,
         docId: doc.id,
@@ -594,19 +631,13 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
 
         this.updateContext({ status: 'success' });
 
-        if (!this.chatContextValue.chatSessionId) {
-          this.updateContext({
-            chatSessionId: AIProvider.LAST_ROOT_SESSION_ID,
-          });
-        }
-
         const { items } = this.chatContextValue;
         const last = items[items.length - 1] as ChatMessage;
         if (!last.id) {
           const historyIds = await AIProvider.histories?.ids(
             doc.workspace.id,
             doc.id,
-            { sessionId: this.chatContextValue.chatSessionId }
+            { sessionId: this.chatSessionId }
           );
           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

@@ -127,6 +127,9 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
   @property({ attribute: false })
   accessor chatContextValue!: ChatContextValue;
 
+  @property({ attribute: false })
+  accessor chatSessionId!: string | undefined;
+
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<ChatContextValue>) => void;
 
@@ -397,8 +400,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
   retry = async () => {
     const { doc } = this.host;
     try {
-      const { chatSessionId } = this.chatContextValue;
+      if (!this.chatSessionId) return;
-      if (!chatSessionId) return;
 
       const abortController = new AbortController();
       const items = [...this.chatContextValue.items];
@@ -410,7 +412,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
       this.updateContext({ items, status: 'loading', error: null });
 
       const stream = AIProvider.actions.chat?.({
-        sessionId: chatSessionId,
+        sessionId: this.chatSessionId,
         retry: true,
         docId: doc.id,
         workspaceId: doc.workspace.id,
@@ -441,7 +443,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
   };
 
   renderEditorActions(item: ChatMessage, isLast: boolean) {
-    const { status, chatSessionId } = this.chatContextValue;
+    const { status } = this.chatContextValue;
 
     if (item.role !== 'assistant') return nothing;
 
@@ -465,7 +467,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
         .actions=${actions}
         .content=${content}
         .isLast=${isLast}
-        .chatSessionId=${chatSessionId ?? undefined}
+        .chatSessionId=${this.chatSessionId}
         .messageId=${messageId}
         .withMargin=${true}
         .retry=${() => this.retry()}
@@ -475,7 +477,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
             .actions=${actions}
             .host=${host}
             .content=${content}
-            .chatSessionId=${chatSessionId ?? undefined}
+            .chatSessionId=${this.chatSessionId}
             .messageId=${messageId ?? undefined}
             .withMargin=${true}
           ></chat-action-list>`

packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/add-popover.ts

@@ -159,7 +159,7 @@ export class ChatPanelAddPopover extends SignalWatcher(
   private readonly _addDocChip = (meta: DocMeta) => {
     this.addChip({
       docId: meta.id,
-      state: 'embedding',
+      state: 'processing',
     });
     this.abortController.abort();
   };

packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts

@@ -24,6 +24,9 @@ export class ChatPanelDocChip extends SignalWatcher(
   @property({ attribute: false })
   accessor chip!: DocChip;
 
+  @property({ attribute: false })
+  accessor addChip!: (chip: ChatChip) => void;
+
   @property({ attribute: false })
   accessor updateChip!: (chip: ChatChip, options: Partial<BaseChip>) => void;
 
@@ -61,9 +64,9 @@ export class ChatPanelDocChip extends SignalWatcher(
     if (
       changedProperties.has('chip') &&
       changedProperties.get('chip')?.state === 'candidate' &&
-      this.chip.state === 'embedding'
+      this.chip.state === 'processing'
     ) {
-      this.embedDocChip().catch(console.error);
+      this.processDocChip().catch(console.error);
     }
   }
 
@@ -74,8 +77,9 @@ export class ChatPanelDocChip extends SignalWatcher(
 
   private readonly onChipClick = async () => {
     if (this.chip.state === 'candidate') {
-      this.updateChip(this.chip, {
+      this.addChip({
-        state: 'embedding',
+        ...this.chip,
+        state: 'processing',
       });
     }
   };
@@ -86,11 +90,11 @@ export class ChatPanelDocChip extends SignalWatcher(
 
   private readonly autoUpdateChip = () => {
     if (this.chip.state !== 'candidate') {
-      this.embedDocChip().catch(console.error);
+      this.processDocChip().catch(console.error);
     }
   };
 
-  private readonly embedDocChip = async () => {
+  private readonly processDocChip = async () => {
     try {
       const doc = this.docDisplayConfig.getDoc(this.chip.docId);
       if (!doc) {
@@ -111,14 +115,14 @@ export class ChatPanelDocChip extends SignalWatcher(
     } catch (e) {
       this.updateChip(this.chip, {
         state: 'failed',
-        tooltip: e instanceof Error ? e.message : 'Failed to embed document',
+        tooltip: e instanceof Error ? e.message : 'Failed to process document',
       });
     }
   };
 
   override render() {
     const { state, docId } = this.chip;
-    const isLoading = state === 'embedding' || state === 'uploading';
+    const isLoading = state === 'processing';
     const getIcon = this.docDisplayConfig.getIcon(docId);
     const docIcon = typeof getIcon === 'function' ? getIcon() : getIcon;
     const icon = getChipIcon(state, docIcon);

packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/file-chip.ts

@@ -15,7 +15,7 @@ export class ChatPanelFileChip extends SignalWatcher(
 
   override render() {
     const { state, fileName, fileType } = this.chip;
-    const isLoading = state === 'embedding' || state === 'uploading';
+    const isLoading = state === 'processing';
     const tooltip = getChipTooltip(state, fileName, this.chip.tooltip);
     const fileIcon = getAttachmentFileIcon(fileType);
     const icon = getChipIcon(state, fileIcon);

packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/utils.ts

@@ -1,3 +1,4 @@
+import type { CopilotContextDoc, CopilotContextFile } from '@affine/graphql';
 import { WarningIcon } from '@blocksuite/icons/lit';
 import { type TemplateResult } from 'lit';
 
@@ -15,14 +16,11 @@ export function getChipTooltip(
   if (state === 'candidate') {
     return 'Click to add doc';
   }
-  if (state === 'embedding') {
+  if (state === 'processing') {
-    return 'Embedding...';
+    return 'Processing...';
-  }
-  if (state === 'uploading') {
-    return 'Uploading...';
   }
   if (state === 'failed') {
-    return 'Failed to embed';
+    return 'Failed to process';
   }
   return name;
 }
@@ -31,7 +29,7 @@ export function getChipIcon(
   state: ChipState,
   icon: TemplateResult<1>
 ): TemplateResult<1> {
-  const isLoading = state === 'embedding' || state === 'uploading';
+  const isLoading = state === 'processing';
   const isFailed = state === 'failed';
   if (isFailed) {
     return WarningIcon();
@@ -50,6 +48,18 @@ export function isFileChip(chip: ChatChip): chip is FileChip {
   return 'fileId' in chip;
 }
 
+export function isDocContext(
+  context: CopilotContextDoc | CopilotContextFile
+): context is CopilotContextDoc {
+  return !('blobId' in context);
+}
+
+export function isFileContext(
+  context: CopilotContextDoc | CopilotContextFile
+): context is CopilotContextFile {
+  return 'blobId' in context;
+}
+
 export function getChipKey(chip: ChatChip) {
   if (isDocChip(chip)) {
     return chip.docId;

packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts

@@ -7,7 +7,7 @@ import {
   NotificationProvider,
   type SpecBuilder,
 } from '@blocksuite/affine/blocks';
-import { debounce, WithDisposable } from '@blocksuite/affine/global/utils';
+import { WithDisposable } from '@blocksuite/affine/global/utils';
 import type { Store } from '@blocksuite/affine/store';
 import { css, html, type PropertyValues } from 'lit';
 import { property, state } from 'lit/decorators.js';
@@ -31,8 +31,10 @@ import type {
   ChatContextValue,
   ChatItem,
   DocChip,
+  FileChip,
 } from './chat-context';
 import type { ChatPanelMessages } from './chat-panel-messages';
+import { isDocContext } from './components/utils';
 
 export class ChatPanel extends WithDisposable(ShadowlessElement) {
   static override styles = css`
@@ -113,52 +115,89 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
   private readonly _chatMessages: Ref<ChatPanelMessages> =
     createRef<ChatPanelMessages>();
 
-  private _resettingCounter = 0;
+  // request counter to track the latest request
+  private _updateHistoryCounter = 0;
 
-  private readonly _resetItems = debounce(() => {
+  private readonly _updateHistory = async () => {
-    const counter = ++this._resettingCounter;
+    const { doc } = this;
     this.isLoading = true;
-    (async () => {
-      const { doc } = this;
 
-      const [histories, actions] = await Promise.all([
+    const currentRequest = ++this._updateHistoryCounter;
-        AIProvider.histories?.chats(doc.workspace.id, doc.id, { fork: false }),
-        AIProvider.histories?.actions(doc.workspace.id, doc.id),
-      ]);
 
-      if (counter !== this._resettingCounter) return;
+    const [histories, actions] = await Promise.all([
+      AIProvider.histories?.chats(doc.workspace.id, doc.id, { fork: false }),
+      AIProvider.histories?.actions(doc.workspace.id, doc.id),
+    ]);
 
-      const items: ChatItem[] = actions ? [...actions] : [];
+    // Check if this is still the latest request
+    if (currentRequest !== this._updateHistoryCounter) {
+      return;
+    }
 
-      if (histories?.at(-1)) {
+    const items: ChatItem[] = actions ? [...actions] : [];
-        const history = histories.at(-1);
-        if (!history) return;
-        this.chatContextValue.chatSessionId = history.sessionId;
-        items.push(...history.messages);
-        AIProvider.LAST_ROOT_SESSION_ID = history.sessionId;
-      }
 
-      const { chips } = this.chatContextValue;
+    if (histories?.at(-1)) {
-      const defaultChip: DocChip = {
+      const history = histories.at(-1);
-        docId: this.doc.id,
+      if (!history) return;
-        state: 'candidate',
+      items.push(...history.messages);
-      };
+      AIProvider.LAST_ROOT_SESSION_ID = history.sessionId;
-      const nextChips =
+    }
-        items.length === 0 && chips.length === 0 ? [defaultChip] : chips;
-      this.chatContextValue = {
-        ...this.chatContextValue,
-        items: items.sort((a, b) => {
-          return (
-            new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
-          );
-        }),
-        chips: nextChips,
-      };
 
-      this.isLoading = false;
+    this.chatContextValue = {
-      this._scrollToEnd();
+      ...this.chatContextValue,
-    })().catch(console.error);
+      items: items.sort(
-  }, 200);
+        (a, b) =>
+          new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
+      ),
+    };
+
+    this.isLoading = false;
+    this._scrollToEnd();
+  };
+
+  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;
+      });
+    }
+
+    this.chatContextValue = {
+      ...this.chatContextValue,
+      chips: chips.length === 0 ? [candidateChip] : chips,
+    };
+  };
 
   @property({ attribute: false })
   accessor host!: EditorHost;
@@ -191,9 +230,12 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
     status: 'idle',
     error: null,
     markdown: '',
-    chatSessionId: null,
   };
 
+  private _chatSessionId: string | undefined;
+
+  private _chatContextId: string | undefined;
+
   private readonly _scrollToEnd = () => {
     this._chatMessages.value?.scrollToEnd();
   };
@@ -214,26 +256,40 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
       })
     ) {
       await AIProvider.histories?.cleanup(this.doc.workspace.id, this.doc.id, [
-        this.chatContextValue.chatSessionId ?? '',
+        this._chatSessionId ?? '',
         ...(
           this.chatContextValue.items.filter(
             item => 'sessionId' in item
           ) as ChatAction[]
         ).map(item => item.sessionId),
       ]);
-      this.chatContextValue.chatSessionId = null;
       notification.toast('History cleared');
-      this._resetItems();
+      await this._updateHistory();
     }
   };
 
+  private readonly _initPanel = async () => {
+    const userId = (await AIProvider.userInfo)?.id;
+    if (!userId) return;
+
+    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
+      );
+    }
+    await this._updateHistory();
+    await this._updateChips();
+  };
+
   protected override updated(_changedProperties: PropertyValues) {
     if (_changedProperties.has('doc')) {
-      requestAnimationFrame(() => {
+      requestAnimationFrame(async () => {
-        this.chatContextValue.chatSessionId = null;
+        await this._initPanel();
-        // TODO get from CopilotContext
-        this.chatContextValue.chips = [];
-        this._resetItems();
       });
     }
 
@@ -266,15 +322,13 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
           event === 'finished' &&
           (status === 'idle' || status === 'success')
         ) {
-          this._resetItems();
+          this._updateHistory().catch(console.error);
         }
       })
     );
     this._disposables.add(
-      AIProvider.slots.userInfo.on(userInfo => {
+      AIProvider.slots.userInfo.on(async () => {
-        if (userInfo) {
+        await this._initPanel();
-          this._resetItems();
-        }
       })
     );
     this._disposables.add(
@@ -318,6 +372,7 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
       <chat-panel-messages
         ${ref(this._chatMessages)}
         .chatContextValue=${this.chatContextValue}
+        .chatSessionId=${this._chatSessionId}
         .updateContext=${this.updateContext}
         .host=${this.host}
         .isLoading=${this.isLoading}
@@ -326,12 +381,14 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
       <chat-panel-chips
         .host=${this.host}
         .chatContextValue=${this.chatContextValue}
+        .chatContextId=${this._chatContextId}
         .updateContext=${this.updateContext}
         .docDisplayConfig=${this.docDisplayConfig}
         .docSearchMenuConfig=${this.docSearchMenuConfig}
       ></chat-panel-chips>
       <chat-panel-input
         .chatContextValue=${this.chatContextValue}
+        .chatSessionId=${this._chatSessionId}
         .networkSearchConfig=${this.networkSearchConfig}
         .updateContext=${this.updateContext}
         .host=${this.host}

packages/frontend/core/src/blocksuite/presets/ai/provider.ts

@@ -74,6 +74,14 @@ export class AIProvider {
     return AIProvider.instance.histories;
   }
 
+  static get session() {
+    return AIProvider.instance.session;
+  }
+
+  static get context() {
+    return AIProvider.instance.context;
+  }
+
   static get actionHistory() {
     return AIProvider.instance.actionHistory;
   }
@@ -100,6 +108,10 @@ export class AIProvider {
 
   private histories: BlockSuitePresets.AIHistoryService | null = null;
 
+  private session: BlockSuitePresets.AISessionService | null = null;
+
+  private context: BlockSuitePresets.AIContextService | null = null;
+
   private toggleGeneralAIOnboarding: ((value: boolean) => void) | null = null;
 
   private forkChat:
@@ -259,6 +271,16 @@ export class AIProvider {
     fn: () => AIUserInfo | Promise<AIUserInfo> | null
   ): void;
 
+  static provide(
+    id: 'session',
+    service: BlockSuitePresets.AISessionService
+  ): void;
+
+  static provide(
+    id: 'context',
+    service: BlockSuitePresets.AIContextService
+  ): void;
+
   static provide(
     id: 'histories',
     service: BlockSuitePresets.AIHistoryService
@@ -292,6 +314,12 @@ export class AIProvider {
     } else if (id === 'histories') {
       AIProvider.instance.histories =
         action as BlockSuitePresets.AIHistoryService;
+    } else if (id === 'session') {
+      AIProvider.instance.session =
+        action as BlockSuitePresets.AISessionService;
+    } else if (id === 'context') {
+      AIProvider.instance.context =
+        action as BlockSuitePresets.AIContextService;
     } else if (id === 'photoEngine') {
       AIProvider.instance.photoEngine =
         action as BlockSuitePresets.AIPhotoEngineService;

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts

@@ -1,6 +1,8 @@
 import { showAILoginRequiredAtom } from '@affine/core/components/affine/auth/ai-login-required';
 import {
+  addContextDocMutation,
   cleanupCopilotSessionMutation,
+  createCopilotContextMutation,
   createCopilotMessageMutation,
   createCopilotSessionMutation,
   forkCopilotSessionMutation,
@@ -9,8 +11,11 @@ import {
   getCopilotSessionsQuery,
   GraphQLError,
   type GraphQLQuery,
+  listContextDocsAndFilesQuery,
+  listContextQuery,
   type QueryOptions,
   type QueryResponse,
+  removeContextDocMutation,
   type RequestOptions,
   updateCopilotSessionMutation,
   UserFriendlyError,
@@ -209,6 +214,74 @@ export class CopilotClient {
     }
   }
 
+  async createContext(workspaceId: string, sessionId: string) {
+    const res = await this.gql({
+      query: createCopilotContextMutation,
+      variables: {
+        workspaceId,
+        sessionId,
+      },
+    });
+    return res.createCopilotContext;
+  }
+
+  async getContextId(workspaceId: string, sessionId: string) {
+    const res = await this.gql({
+      query: listContextQuery,
+      variables: {
+        workspaceId,
+        sessionId,
+      },
+    });
+    return res.currentUser?.copilot?.contexts?.[0]?.id;
+  }
+
+  async addContextDoc(options: OptionsField<typeof addContextDocMutation>) {
+    const res = await this.gql({
+      query: addContextDocMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.addContextDoc;
+  }
+
+  async removeContextDoc(
+    options: OptionsField<typeof removeContextDocMutation>
+  ) {
+    const res = await this.gql({
+      query: removeContextDocMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.removeContextDoc;
+  }
+
+  async addContextFile() {
+    return;
+  }
+
+  async removeContextFile() {
+    return;
+  }
+
+  async getContextDocsAndFiles(
+    workspaceId: string,
+    sessionId: string,
+    contextId: string
+  ) {
+    const res = await this.gql({
+      query: listContextDocsAndFilesQuery,
+      variables: {
+        workspaceId,
+        sessionId,
+        contextId,
+      },
+    });
+    return res.currentUser?.copilot?.contexts?.[0];
+  }
+
   async chatText({
     sessionId,
     messageId,

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts

@@ -31,22 +31,29 @@ export type ToImageOptions = TextToTextOptions & {
   seed?: string;
 };
 
-export function createChatSession({
+export async function createChatSession({
   client,
   workspaceId,
   docId,
-  promptName,
+  promptName = 'Chat With AFFiNE AI',
 }: {
   client: CopilotClient;
   workspaceId: string;
   docId: string;
-  promptName: string;
+  promptName?: string;
 }) {
-  return client.createSession({
+  const sessionId = await client.createSession({
     workspaceId,
     docId,
     promptName,
   });
+  // always update the prompt name
+  await updateChatSession({
+    sessionId,
+    client,
+    promptName,
+  });
+  return sessionId;
 }
 
 export function updateChatSession({
@@ -119,7 +126,8 @@ async function createSessionMessage({
   }
   const hasAttachments = attachments && attachments.length > 0;
   const sessionId = await (providedSessionId ??
-    client.createSession({
+    createChatSession({
+      client,
       workspaceId,
       docId,
       promptName: promptName as string,

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx

@@ -1,12 +1,10 @@
 import { AIProvider } from '@affine/core/blocksuite/presets/ai';
 import { toggleGeneralAIOnboarding } from '@affine/core/components/affine/ai-onboarding/apis';
-import type { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
 import type { GlobalDialogService } from '@affine/core/modules/dialogs';
 import {
   type getCopilotHistoriesQuery,
   type RequestOptions,
 } from '@affine/graphql';
-import { UnauthorizedError } from '@blocksuite/affine/blocks';
 import { assertExists } from '@blocksuite/affine/global/utils';
 import { z } from 'zod';
 
@@ -39,80 +37,19 @@ const processTypeToPromptName = new Map(
   })
 );
 
-// a single workspace should have only a single chat session
-// user-id:workspace-id:doc-id -> chat session id
-const chatSessions = new Map<
-  string,
-  { getSessionId: Promise<string>; promptName: string }
->();
-
 export function setupAIProvider(
   client: CopilotClient,
-  globalDialogService: GlobalDialogService,
+  globalDialogService: GlobalDialogService
-  networkSearchService: AINetworkSearchService
 ) {
-  function getChatPrompt(options: BlockSuitePresets.ChatOptions) {
-    const { attachments, docs } = options;
-    if (attachments?.length || docs?.length) {
-      return 'Chat With AFFiNE AI';
-    }
-    const { enabled, visible } = networkSearchService;
-    return visible.value && enabled.value
-      ? 'Search With AFFiNE AI'
-      : 'Chat With AFFiNE AI';
-  }
-  async function getChatSessionId(options: BlockSuitePresets.ChatOptions) {
-    const userId = (await AIProvider.userInfo)?.id;
-
-    if (!userId) {
-      throw new UnauthorizedError();
-    }
-
-    const { workspaceId, docId } = options;
-    const storeKey = `${userId}:${workspaceId}:${docId}`;
-    const promptName = getChatPrompt(options);
-    if (!chatSessions.has(storeKey)) {
-      chatSessions.set(storeKey, {
-        getSessionId: createChatSession({
-          client,
-          workspaceId,
-          docId,
-          promptName,
-        }).then(sessionId => {
-          return updateChatSession({
-            sessionId,
-            client,
-            promptName,
-          });
-        }),
-        promptName,
-      });
-    }
-    try {
-      /* oxlint-disable @typescript-eslint/no-non-null-assertion */
-      const { getSessionId, promptName: prevName } =
-        chatSessions.get(storeKey)!;
-      const sessionId = await getSessionId;
-      //update prompt name
-      if (prevName !== promptName) {
-        await updateChatSession({
-          sessionId,
-          client,
-          promptName,
-        });
-        chatSessions.set(storeKey, { getSessionId, promptName });
-      }
-      return sessionId;
-    } catch (err) {
-      // do not cache the error
-      chatSessions.delete(storeKey);
-      throw err;
-    }
-  }
-
   //#region actions
   AIProvider.provide('chat', options => {
-    const sessionId = options.sessionId ?? getChatSessionId(options);
+    const sessionId =
+      options.sessionId ??
+      createChatSession({
+        client,
+        workspaceId: options.workspaceId,
+        docId: options.docId,
+      });
     const { input, docs, ...rest } = options;
     const params = docs?.length
       ? {
@@ -473,6 +410,56 @@ Could you make a new website based on these notes and send back just the html fi
   });
   //#endregion
 
+  AIProvider.provide('session', {
+    createSession: async (
+      workspaceId: string,
+      docId: string,
+      promptName?: string
+    ) => {
+      return createChatSession({
+        client,
+        workspaceId,
+        docId,
+        promptName,
+      });
+    },
+    updateSession: async (sessionId: string, promptName: string) => {
+      return updateChatSession({
+        client,
+        sessionId,
+        promptName,
+      });
+    },
+  });
+
+  AIProvider.provide('context', {
+    createContext: async (workspaceId: string, sessionId: string) => {
+      return client.createContext(workspaceId, sessionId);
+    },
+    getContextId: async (workspaceId: string, sessionId: string) => {
+      return client.getContextId(workspaceId, sessionId);
+    },
+    addContextDoc: async (options: { contextId: string; docId: string }) => {
+      return client.addContextDoc(options);
+    },
+    removeContextDoc: async (options: { contextId: string; docId: string }) => {
+      return client.removeContextDoc(options);
+    },
+    addContextFile: async () => {
+      return client.addContextFile();
+    },
+    removeContextFile: async () => {
+      return client.removeContextFile();
+    },
+    getContextDocsAndFiles: async (
+      workspaceId: string,
+      sessionId: string,
+      contextId: string
+    ) => {
+      return client.getContextDocsAndFiles(workspaceId, sessionId, contextId);
+    },
+  });
+
   AIProvider.provide('histories', {
     actions: async (
       workspaceId: string,

packages/frontend/core/src/components/providers/workspace-side-effects.tsx

@@ -8,7 +8,6 @@ import { SyncAwareness } from '@affine/core/components/affine/awareness';
 import { useRegisterFindInPageCommands } from '@affine/core/components/hooks/affine/use-register-find-in-page-commands';
 import { useRegisterWorkspaceCommands } from '@affine/core/components/hooks/use-register-workspace-commands';
 import { OverCapacityNotification } from '@affine/core/components/over-capacity';
-import { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
 import {
   EventSourceService,
   FetchService,
@@ -144,7 +143,6 @@ export const WorkspaceSideEffects = () => {
   const graphqlService = useService(GraphQLService);
   const eventSourceService = useService(EventSourceService);
   const fetchService = useService(FetchService);
-  const networkSearchService = useService(AINetworkSearchService);
 
   useEffect(() => {
     const dispose = setupAIProvider(
@@ -153,8 +151,7 @@ export const WorkspaceSideEffects = () => {
         fetchService.fetch,
         eventSourceService.eventSource
       ),
-      globalDialogService,
+      globalDialogService
-      networkSearchService
     );
     return () => {
       dispose();
@@ -164,7 +161,6 @@ export const WorkspaceSideEffects = () => {
     fetchService,
     workspaceDialogService,
     graphqlService,
-    networkSearchService,
     globalDialogService,
   ]);
 

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

@@ -193,11 +193,31 @@ export const removeContextDocMutation = {
   definitionName: 'removeContextDoc',
   containsFile: false,
   query: `
-mutation removeContextDoc($options: RemoveContextFileInput!) {
+mutation removeContextDoc($options: RemoveContextDocInput!) {
   removeContextDoc(options: $options)
 }`,
 };
 
+export const listContextDocsAndFilesQuery = {
+  id: 'listContextDocsAndFilesQuery' as const,
+  operationName: 'listContextDocsAndFiles',
+  definitionName: 'currentUser',
+  containsFile: false,
+  query: `
+query listContextDocsAndFiles($workspaceId: String!, $sessionId: String!, $contextId: String!) {
+  currentUser {
+    copilot(workspaceId: $workspaceId) {
+      contexts(sessionId: $sessionId, contextId: $contextId) {
+        docs {
+          id
+          createdAt
+        }
+      }
+    }
+  }
+}`,
+};
+
 export const listContextQuery = {
   id: 'listContextQuery' as const,
   operationName: 'listContext',

packages/frontend/graphql/src/schema.ts

@@ -42,6 +42,11 @@ export interface AddContextDocInput {
   docId: Scalars['String']['input'];
 }
 
+export interface RemoveContextDocInput {
+  contextId: Scalars['String']['input'];
+  docId: Scalars['String']['input'];
+}
+
 export interface AlreadyInSpaceDataType {
   __typename?: 'AlreadyInSpaceDataType';
   spaceId: Scalars['String']['output'];
@@ -1013,7 +1018,7 @@ export interface MutationReleaseDeletedBlobsArgs {
 }
 
 export interface MutationRemoveContextDocArgs {
-  options: RemoveContextFileInput;
+  options: RemoveContextDocInput;
 }
 
 export interface MutationRemoveWorkspaceFeatureArgs {
@@ -1310,11 +1315,6 @@ export interface RemoveAvatar {
   success: Scalars['Boolean']['output'];
 }
 
-export interface RemoveContextFileInput {
-  contextId: Scalars['String']['input'];
-  fileId: Scalars['String']['input'];
-}
-
 export interface RevokeDocUserRoleInput {
   docId: Scalars['String']['input'];
   userId: Scalars['String']['input'];
@@ -1952,28 +1952,56 @@ export type AddContextDocMutationVariables = Exact<{
   options: AddContextDocInput;
 }>;
 
+export type RemoveContextDocMutationVariables = Exact<{
+  options: RemoveContextDocInput;
+}>;
+
 export type AddContextDocMutation = {
   __typename?: 'Mutation';
   addContextDoc: Array<{
     __typename?: 'CopilotContextListItem';
     id: string;
-    createdAt: number;
-    name: string | null;
-    chunkSize: number | null;
-    status: ContextFileStatus | null;
-    blobId: string | null;
   }>;
 };
 
-export type RemoveContextDocMutationVariables = Exact<{
-  options: RemoveContextFileInput;
-}>;
-
 export type RemoveContextDocMutation = {
   __typename?: 'Mutation';
   removeContextDoc: boolean;
 };
 
+export type ListContextDocsAndFilesQueryVariables = Exact<{
+  workspaceId: Scalars['String']['input'];
+  sessionId: Scalars['String']['input'];
+  contextId: Scalars['String']['input'];
+}>;
+
+export type ListContextDocsAndFilesQuery = {
+  __typename?: 'Query';
+  currentUser: {
+    __typename?: 'UserType';
+    copilot: {
+      __typename?: 'Copilot';
+      contexts: Array<{
+        __typename?: 'CopilotContext';
+        docs: Array<{
+          __typename?: 'CopilotContextDoc';
+          id: string;
+          createdAt: number;
+        }>;
+        files: Array<{
+          __typename?: 'CopilotContextFile';
+          id: string;
+          name: string;
+          blobId: string;
+          chunkSize: number;
+          status: ContextFileStatus;
+          createdAt: number;
+        }>;
+      }>;
+    };
+  } | null;
+};
+
 export type ListContextQueryVariables = Exact<{
   workspaceId: Scalars['String']['input'];
   sessionId: Scalars['String']['input'];
@@ -3372,6 +3400,11 @@ export type Queries =
       variables: ListBlobsQueryVariables;
       response: ListBlobsQuery;
     }
+  | {
+      name: 'listContextDocsAndFilesQuery';
+      variables: ListContextDocsAndFilesQueryVariables;
+      response: ListContextDocsAndFilesQuery;
+    }
   | {
       name: 'listContextQuery';
       variables: ListContextQueryVariables;

tests/affine-cloud-copilot/e2e/copilot.spec.ts

@@ -10,10 +10,7 @@ import {
   getBlockSuiteEditorTitle,
   waitForEditorLoad,
 } from '@affine-test/kit/utils/page-logic';
-import {
+import { clickSideBarAllPageButton } from '@affine-test/kit/utils/sidebar';
-  clickSideBarAllPageButton,
-  clickSideBarUseAvatar,
-} from '@affine-test/kit/utils/sidebar';
 import { createLocalWorkspace } from '@affine-test/kit/utils/workspace';
 import { expect, type Page } from '@playwright/test';
 
@@ -394,20 +391,12 @@ test.describe('chat panel', () => {
     await page.waitForTimeout(200);
     await createLocalWorkspace({ name: 'test' }, page);
     await clickNewPageButton(page);
-    await clickSideBarUseAvatar(page);
+
-    await page.getByTestId('workspace-modal-account-settings-option').click();
-    await page.getByTestId('experimental-features-trigger').click();
-    await page
-      .getByTestId('experimental-prompt')
-      .getByTestId('affine-checkbox')
-      .click();
-    await page.getByTestId('experimental-confirm-button').click();
-    await page.getByTestId('enable_ai_network_search').click();
-    await page.getByTestId('modal-close-button').click();
     await openChat(page);
     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',
@@ -423,6 +412,7 @@ 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',
@@ -770,11 +760,13 @@ test.describe('chat with doc', () => {
     // oxlint-disable-next-line unicorn/prefer-dom-node-dataset
     expect(await chip.getAttribute('data-state')).toBe('candidate');
     await chip.click();
+    await page.waitForTimeout(1000);
     // oxlint-disable-next-line unicorn/prefer-dom-node-dataset
     expect(await chip.getAttribute('data-state')).toBe('success');
 
     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',
@@ -786,5 +778,13 @@ test.describe('chat with doc', () => {
     ).toBeGreaterThan(0);
     await clearChat(page);
     expect((await collectChat(page)).length).toBe(0);
+
+    await page.reload();
+    await page.waitForTimeout(1000);
+    await openChat(page);
+    expect(await chipTitle.textContent()).toBe('AFFiNE AI');
+    const chip2 = await page.getByTestId('chat-panel-chip');
+    // oxlint-disable-next-line unicorn/prefer-dom-node-dataset
+    expect(await chip2.getAttribute('data-state')).toBe('success');
   });
 });