feat(core): add ai file context api (#10842)

Close [BS-2349](https://linear.app/affine-design/issue/BS-2349).

### What Changed?
- Add file context graphql apis
- Pass matched file chunks to LLM

[录屏2025-02-19 23.27.47.mov <span class="graphite__hidden">(uploaded via Graphite)</span> <img class="graphite__hidden" src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/8e8a98ca-6959-4bb6-9759-b51d97cede49.mov" />](https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/8e8a98ca-6959-4bb6-9759-b51d97cede49.mov)

blocksuite/affine/shared/src/services/file-size-limit-service.ts

@@ -1,7 +1,7 @@
 import { StoreExtension } from '@blocksuite/store';
 
 // bytes.parse('2GB')
-const maxFileSize = 2147483648;
+const maxFileSize = 2 * 1024 * 1024 * 1024;
 
 export class FileSizeLimitService extends StoreExtension {
   static override key = 'file-size-limit';

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

@@ -141,6 +141,7 @@ export class ChatSession implements AsyncDisposable {
   finish(params: PromptParams): PromptMessage[] {
     const messages = this.takeMessages();
     const firstMessage = messages.at(0);
+    // TODO: refactor this {{content}} keyword agreement
     // if the message in prompt config contains {{content}},
     // we should combine it with the user message in the prompt
     if (

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

@@ -1,5 +1,6 @@
 import type {
   ChatHistoryOrder,
+  ContextMatchedFileChunk,
   CopilotContextDoc,
   CopilotContextFile,
   CopilotSessionType,
@@ -10,7 +11,7 @@ import type { EditorHost } from '@blocksuite/affine/block-std';
 import type { GfxModel } from '@blocksuite/affine/block-std/gfx';
 import type { BlockModel } from '@blocksuite/affine/store';
 
-import type { DocContext } from '../chat-panel/chat-context';
+import type { DocContext, FileContext } from '../chat-panel/chat-context';
 
 export const translateLangs = [
   'English',
@@ -114,7 +115,10 @@ declare global {
     interface ChatOptions extends AITextActionOptions {
       sessionId?: string;
       isRootSession?: boolean;
-      docs?: DocContext[];
+      contexts?: {
+        docs: DocContext[];
+        files: FileContext[];
+      };
     }
 
     interface TranslateOptions extends AITextActionOptions {
@@ -250,19 +254,22 @@ declare global {
       addContextDoc: (options: {
         contextId: string;
         docId: string;
-      }) => Promise<{ id: string; createdAt: number }>;
+      }) => Promise<CopilotContextDoc>;
       removeContextDoc: (options: {
         contextId: string;
         docId: string;
       }) => Promise<boolean>;
-      addContextFile: (options: {
+      addContextFile: (
+        file: File,
+        options: {
           contextId: string;
-        fileId: string;
+          blobId: string;
-      }) => Promise<void>;
+        }
+      ) => Promise<CopilotContextFile>;
       removeContextFile: (options: {
         contextId: string;
         fileId: string;
-      }) => Promise<void>;
+      }) => Promise<boolean>;
       getContextDocsAndFiles: (
         workspaceId: string,
         sessionId: string,
@@ -274,6 +281,11 @@ declare global {
           }
         | undefined
       >;
+      matchContext: (
+        contextId: string,
+        content: string,
+        limit?: number
+      ) => Promise<ContextMatchedFileChunk[] | undefined>;
     }
 
     // TODO(@Peng): should be refactored to get rid of implement details (like messages, action, role, etc.)

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

@@ -36,11 +36,18 @@ export type ChatStatus =
 
 export interface DocContext {
   docId: string;
-  plaintext?: string;
+  refIndex: number;
-  markdown?: string;
+  markdown: string;
-  images?: File[];
 }
 
+export type FileContext = {
+  blobId: string;
+  refIndex: number;
+  fileName: string;
+  fileType: string;
+  chunks: string;
+};
+
 export type ChatContextValue = {
   // history messages of the chat
   items: ChatItem[];
@@ -73,19 +80,19 @@ export interface BaseChip {
    * failed: the chip is failed to process
    */
   state: ChipState;
-  tooltip?: string;
+  tooltip?: string | null;
 }
 
 export interface DocChip extends BaseChip {
   docId: string;
-  markdown?: Signal<string>;
+  markdown?: Signal<string> | null;
-  tokenCount?: number;
+  tokenCount?: number | null;
 }
 
 export interface FileChip extends BaseChip {
-  fileName: string;
+  file: File;
-  fileId: string;
+  fileId?: string | null;
-  fileType: string;
+  blobId?: string | null;
 }
 
 export interface TagChip extends BaseChip {

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

@@ -112,6 +112,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
           if (isFileChip(chip)) {
             return html`<chat-panel-file-chip
               .chip=${chip}
+              .removeChip=${this._removeChip}
             ></chat-panel-file-chip>`;
           }
           return null;
@@ -174,14 +175,15 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
   };
 
   private readonly _addChip = async (chip: ChatChip) => {
+    this.isCollapsed = false;
     if (
       this.chatContextValue.chips.length === 1 &&
       this.chatContextValue.chips[0].state === 'candidate'
     ) {
-      await this._addToContext(chip);
       this.updateContext({
         chips: [chip],
       });
+      await this._addToContext(chip);
       return;
     }
     // remove the chip if it already exists
@@ -189,16 +191,16 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
       if (isDocChip(chip)) {
         return !isDocChip(item) || item.docId !== chip.docId;
       } else {
-        return !isFileChip(item) || item.fileId !== chip.fileId;
+        return !isFileChip(item) || item.file !== chip.file;
       }
     });
+    this.updateContext({
+      chips: [...chips, chip],
+    });
     if (chips.length < this.chatContextValue.chips.length) {
       await this._removeFromContext(chip);
     }
     await this._addToContext(chip);
-    this.updateContext({
-      chips: [...chips, chip],
-    });
   };
 
   private readonly _updateChip = (
@@ -209,7 +211,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
       if (isDocChip(chip)) {
         return isDocChip(item) && item.docId === chip.docId;
       } else {
-        return isFileChip(item) && item.fileId === chip.fileId;
+        return isFileChip(item) && item.file === chip.file;
       }
     });
     const nextChip: ChatChip = {
@@ -237,7 +239,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
       await this._removeFromContext(chip);
       this.updateContext({
         chips: this.chatContextValue.chips.filter(item => {
-          return !isFileChip(item) || item.fileId !== chip.fileId;
+          return !isFileChip(item) || item.file !== chip.file;
         }),
       });
     }
@@ -254,10 +256,23 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
         docId: chip.docId,
       });
     } else {
-      await AIProvider.context.addContextFile({
+      try {
+        const blobId = await this.host.doc.blobSync.set(chip.file);
+        const contextFile = await AIProvider.context.addContextFile(chip.file, {
           contextId,
-        fileId: chip.fileId,
+          blobId,
         });
+        this._updateChip(chip, {
+          state: 'success',
+          blobId: contextFile.blobId,
+          fileId: contextFile.id,
+        });
+      } catch (e) {
+        this._updateChip(chip, {
+          state: 'failed',
+          tooltip: e instanceof Error ? e.message : 'Add context file error',
+        });
+      }
     }
   };
 
@@ -271,7 +286,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
         contextId,
         docId: chip.docId,
       });
-    } else {
+    } else if (isFileChip(chip) && chip.fileId) {
       await AIProvider.context.removeContextFile({
         contextId,
         fileId: chip.fileId,

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

@@ -19,8 +19,13 @@ import { AIProvider } from '../provider';
 import { reportResponse } from '../utils/action-reporter';
 import { readBlobAsURL } from '../utils/image';
 import type { AINetworkSearchConfig } from './chat-config';
-import type { ChatContextValue, ChatMessage, DocContext } from './chat-context';
+import type {
-import { isDocChip } from './components/utils';
+  ChatContextValue,
+  ChatMessage,
+  DocContext,
+  FileContext,
+} from './chat-context';
+import { isDocChip, isFileChip } from './components/utils';
 import { PROMPT_NAME_AFFINE_AI, PROMPT_NAME_NETWORK_SEARCH } from './const';
 
 const MaximumImageCount = 32;
@@ -199,6 +204,9 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
   @property({ attribute: false })
   accessor getSessionId!: () => Promise<string | undefined>;
 
+  @property({ attribute: false })
+  accessor getContextId!: () => Promise<string | undefined>;
+
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<ChatContextValue>) => void;
 
@@ -218,7 +226,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
   private get _isNetworkDisabled() {
     return (
       !!this.chatContextValue.images.length ||
-      !!this.chatContextValue.chips.filter(chip => chip.state !== 'candidate')
+      !!this.chatContextValue.chips.filter(chip => chip.state === 'success')
         .length
     );
   }
@@ -452,7 +460,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
   };
 
   send = async (text: string) => {
-    const { status, markdown, chips, images } = this.chatContextValue;
+    const { status, markdown, images } = this.chatContextValue;
     if (status === 'loading' || status === 'transmitting') return;
     if (!text) return;
 
@@ -498,17 +506,12 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
 
       const abortController = new AbortController();
       const sessionId = await this.getSessionId();
-      const docs: DocContext[] = chips
+
-        .filter(isDocChip)
+      const contexts = await this._getMatchedContexts(userInput);
-        .filter(chip => !!chip.markdown?.value && chip.state === 'success')
-        .map(chip => ({
-          docId: chip.docId,
-          markdown: chip.markdown?.value || '',
-        }));
       const stream = AIProvider.actions.chat?.({
         sessionId,
         input: userInput,
-        docs: docs,
+        contexts,
         docId: doc.id,
         attachments: images,
         workspaceId: doc.workspace.id,
@@ -550,6 +553,44 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
       this.updateContext({ abortController: null });
     }
   };
+
+  private async _getMatchedContexts(userInput: string) {
+    const contextId = await this.getContextId();
+    const matched = contextId
+      ? (await AIProvider.context?.matchContext(contextId, userInput)) || []
+      : [];
+    const contexts = this.chatContextValue.chips.reduce(
+      (acc, chip, index) => {
+        if (chip.state !== 'success') {
+          return acc;
+        }
+        if (isDocChip(chip) && !!chip.markdown?.value) {
+          acc.docs.push({
+            docId: chip.docId,
+            refIndex: index + 1,
+            markdown: chip.markdown.value,
+          });
+        }
+        if (isFileChip(chip) && chip.blobId) {
+          const matchedChunks = matched
+            .filter(chunk => chunk.fileId === chip.fileId)
+            .map(chunk => chunk.content);
+          if (matchedChunks.length > 0) {
+            acc.files.push({
+              blobId: chip.blobId,
+              refIndex: index + 1,
+              fileName: chip.file.name,
+              fileType: chip.file.type,
+              chunks: matchedChunks.join('\n'),
+            });
+          }
+        }
+        return acc;
+      },
+      { docs: [], files: [] } as { docs: DocContext[]; files: FileContext[] }
+    );
+    return contexts;
+  }
 }
 
 declare global {

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

@@ -1,8 +1,10 @@
+import { toast } from '@affine/component';
 import { ShadowlessElement } from '@blocksuite/affine/block-std';
-import type { LinkedMenuGroup } from '@blocksuite/affine/blocks/root';
+import { type LinkedMenuGroup } from '@blocksuite/affine/blocks/root';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { scrollbarStyle } from '@blocksuite/affine/shared/styles';
-import { SearchIcon } from '@blocksuite/icons/lit';
+import { openFileOrFiles } from '@blocksuite/affine/shared/utils';
+import { SearchIcon, UploadIcon } from '@blocksuite/icons/lit';
 import type { DocMeta } from '@blocksuite/store';
 import { css, html } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
@@ -138,9 +140,35 @@ export class ChatPanelAddPopover extends SignalWatcher(
             })
           : html`<div class="no-result">No Result</div>`}
       </div>
+      <div class="divider"></div>
+      <div class="upload-wrapper">
+        <icon-button
+          width="280px"
+          height="30px"
+          data-id="upload"
+          .text=${'Upload files (pdf, txt, csv)'}
+          @click=${this._addFileChip}
+        >
+          ${UploadIcon()}
+        </icon-button>
+      </div>
     </div>`;
   }
 
+  private readonly _addFileChip = async () => {
+    const file = await openFileOrFiles();
+    if (!file) return;
+    if (file.size > 50 * 1024 * 1024) {
+      toast('You can only upload files less than 50MB');
+      return;
+    }
+    this.addChip({
+      file,
+      state: 'processing',
+    });
+    this.abortController.abort();
+  };
+
   private _onInput(event: Event) {
     this._query = (event.target as HTMLInputElement).value;
     this._updateDocGroup();

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

@@ -125,7 +125,7 @@ export class ChatPanelDocChip extends SignalWatcher(
     } catch (e) {
       this.updateChip(this.chip, {
         state: 'failed',
-        tooltip: e instanceof Error ? e.message : 'Failed to process document',
+        tooltip: e instanceof Error ? e.message : 'Failed to extract markdown',
       });
     }
   };

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

@@ -4,7 +4,7 @@ import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { html } from 'lit';
 import { property } from 'lit/decorators.js';
 
-import type { FileChip } from '../chat-context';
+import type { ChatChip, FileChip } from '../chat-context';
 import { getChipIcon, getChipTooltip } from './utils';
 
 export class ChatPanelFileChip extends SignalWatcher(
@@ -13,19 +13,28 @@ export class ChatPanelFileChip extends SignalWatcher(
   @property({ attribute: false })
   accessor chip!: FileChip;
 
+  @property({ attribute: false })
+  accessor removeChip!: (chip: ChatChip) => void;
+
   override render() {
-    const { state, fileName, fileType } = this.chip;
+    const { state, file } = this.chip;
     const isLoading = state === 'processing';
-    const tooltip = getChipTooltip(state, fileName, this.chip.tooltip);
+    const tooltip = getChipTooltip(state, file.name, this.chip.tooltip);
+    const fileType = file.name.split('.').pop() ?? '';
     const fileIcon = getAttachmentFileIcon(fileType);
     const icon = getChipIcon(state, fileIcon);
 
     return html`<chat-panel-chip
       .state=${state}
-      .name=${fileName}
+      .name=${file.name}
       .tooltip=${tooltip}
       .icon=${icon}
       .closeable=${!isLoading}
+      .onChipDelete=${this.onChipDelete}
     ></chat-panel-chip>`;
   }
+
+  private readonly onChipDelete = () => {
+    this.removeChip(this.chip);
+  };
 }

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

@@ -8,7 +8,7 @@ import type { ChatChip, ChipState, DocChip, FileChip } from '../chat-context';
 export function getChipTooltip(
   state: ChipState,
   name: string,
-  tooltip?: string
+  tooltip?: string | null
 ) {
   if (tooltip) {
     return tooltip;
@@ -20,7 +20,7 @@ export function getChipTooltip(
     return 'Processing...';
   }
   if (state === 'failed') {
-    return 'Failed to process';
+    return 'Failed to add to context';
   }
   return name;
 }
@@ -45,7 +45,7 @@ export function isDocChip(chip: ChatChip): chip is DocChip {
 }
 
 export function isFileChip(chip: ChatChip): chip is FileChip {
-  return 'fileId' in chip;
+  return 'file' in chip && chip.file instanceof File;
 }
 
 export function isDocContext(

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

@@ -28,6 +28,7 @@ import type {
   DocSearchMenuConfig,
 } from './chat-config';
 import type {
+  ChatChip,
   ChatContextValue,
   ChatItem,
   DocChip,
@@ -180,7 +181,6 @@ export class ChatPanel extends SignalWatcher(
     }
 
     // context initialized, show the chips
-    let chips: (DocChip | FileChip)[] = [];
     const { docs = [], files = [] } =
       (await AIProvider.context?.getContextDocsAndFiles(
         this.doc.workspace.id,
@@ -191,23 +191,34 @@ export class ChatPanel extends SignalWatcher(
       (a, b) =>
         new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
     );
-    chips = list.map(item => {
+    const chips: ChatChip[] = await Promise.all(
-      let chip: DocChip | FileChip;
+      list.map(async item => {
         if (isDocContext(item)) {
-        chip = {
+          return {
             docId: item.id,
             state: 'processing',
-        };
+          } as DocChip;
+        }
+        const file = await this.host.doc.blobSync.get(item.blobId);
+        if (!file) {
+          return {
+            blobId: item.id,
+            file: new File([], item.name),
+            state: 'failed',
+            tooltip: 'File not found in blob storage',
+          } as FileChip;
         } else {
-        chip = {
+          return {
+            file: new File([file], item.name),
+            blobId: item.blobId,
             fileId: item.id,
             state: item.status === 'finished' ? 'success' : item.status,
-          fileName: item.name,
+            tooltip: item.error,
-          fileType: '',
+          } as FileChip;
-        };
         }
-      return chip;
+      })
-    });
+    );
+
     this.chatContextValue = {
       ...this.chatContextValue,
       chips,
@@ -489,6 +500,7 @@ export class ChatPanel extends SignalWatcher(
       <chat-panel-input
         .chatContextValue=${this.chatContextValue}
         .getSessionId=${this._getSessionId}
+        .getContextId=${this._getContextId}
         .networkSearchConfig=${this.networkSearchConfig}
         .updateContext=${this.updateContext}
         .host=${this.host}

packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts

@@ -7,6 +7,7 @@ import { showAILoginRequiredAtom } from '@affine/core/components/affine/auth/ai-
 import type { UserFriendlyError } from '@affine/error';
 import {
   addContextDocMutation,
+  addContextFileMutation,
   cleanupCopilotSessionMutation,
   createCopilotContextMutation,
   createCopilotMessageMutation,
@@ -22,6 +23,7 @@ import {
   type QueryOptions,
   type QueryResponse,
   removeContextDocMutation,
+  removeContextFileMutation,
   type RequestOptions,
   updateCopilotSessionMutation,
 } from '@affine/graphql';
@@ -261,12 +263,30 @@ export class CopilotClient {
     return res.removeContextDoc;
   }
 
-  async addContextFile() {
+  async addContextFile(
-    return;
+    content: File,
+    options: OptionsField<typeof addContextFileMutation>
+  ) {
+    const res = await this.gql({
+      query: addContextFileMutation,
+      variables: {
+        content,
+        options,
+      },
+    });
+    return res.addContextFile;
   }
 
-  async removeContextFile() {
+  async removeContextFile(
-    return;
+    options: OptionsField<typeof removeContextFileMutation>
+  ) {
+    const res = await this.gql({
+      query: removeContextFileMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.removeContextFile;
   }
 
   async getContextDocsAndFiles(

packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx

@@ -36,21 +36,12 @@ export function setupAIProvider(
 ) {
   //#region actions
   AIProvider.provide('chat', options => {
-    const { input, docs, ...rest } = options;
+    const { input, contexts, ...rest } = options;
-    const params = docs?.length
-      ? {
-          docs: docs.map((doc, i) => ({
-            docId: doc.docId,
-            markdown: doc.markdown,
-            index: i + 1,
-          })),
-        }
-      : undefined;
     return textToText({
       ...rest,
       client,
       content: input,
-      params,
+      params: contexts,
     });
   });
 
@@ -441,11 +432,17 @@ Could you make a new website based on these notes and send back just the html fi
     removeContextDoc: async (options: { contextId: string; docId: string }) => {
       return client.removeContextDoc(options);
     },
-    addContextFile: async () => {
+    addContextFile: async (
-      return client.addContextFile();
+      file: File,
+      options: { contextId: string; blobId: string }
+    ) => {
+      return client.addContextFile(file, options);
     },
-    removeContextFile: async () => {
+    removeContextFile: async (options: {
-      return client.removeContextFile();
+      contextId: string;
+      fileId: string;
+    }) => {
+      return client.removeContextFile(options);
     },
     getContextDocsAndFiles: async (
       workspaceId: string,
@@ -454,6 +451,13 @@ Could you make a new website based on these notes and send back just the html fi
     ) => {
       return client.getContextDocsAndFiles(workspaceId, sessionId, contextId);
     },
+    matchContext: async (
+      contextId: string,
+      content: string,
+      limit?: number
+    ) => {
+      return client.matchContext(contextId, content, limit);
+    },
   });
 
   AIProvider.provide('histories', {

packages/frontend/graphql/src/graphql/copilot-context-file-add.gql

@@ -4,6 +4,7 @@ mutation addContextFile($content: Upload!, $options: AddContextFileInput!) {
     createdAt
     name
     chunkSize
+    error
     status
     blobId
   }

packages/frontend/graphql/src/graphql/copilot-context-file-list.gql

@@ -15,6 +15,7 @@ query listContextFiles(
           name
           blobId
           chunkSize
+          error
           status
           createdAt
         }

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

@@ -170,6 +170,7 @@ export const addContextFileMutation = {
     createdAt
     name
     chunkSize
+    error
     status
     blobId
   }
@@ -193,6 +194,7 @@ export const listContextFilesQuery = {
           name
           blobId
           chunkSize
+          error
           status
           createdAt
         }

packages/frontend/graphql/src/schema.ts

@@ -2355,6 +2355,7 @@ export type AddContextFileMutation = {
     createdAt: number;
     name: string;
     chunkSize: number;
+    error: string | null;
     status: ContextFileStatus;
     blobId: string;
   };
@@ -2385,6 +2386,7 @@ export type ListContextFilesQuery = {
           name: string;
           blobId: string;
           chunkSize: number;
+          error: string | null;
           status: ContextFileStatus;
           createdAt: number;
         }>;

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

@@ -1040,7 +1040,7 @@ test.describe('chat with doc', () => {
       .nth(0);
     await richText.click(); // Ensure proper focus
     await page.keyboard.type(
-      'AFFiNE AI is an assistant with the ability to create well-structured outlines for any given content.',
+      'AFiAI is an assistant with the ability to create well-structured outlines for any given content.',
       {
         delay: 50,
       }
@@ -1050,21 +1050,18 @@ test.describe('chat with doc', () => {
     // 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 typeChatSequentially(page, 'What is AFiAI?');
     await page.keyboard.press('Enter');
     const history = await collectChat(page);
     expect(history[0]).toEqual({
       name: 'You',
-      content: 'What is AFFiNE AI?',
+      content: 'What is AFiAI?',
     });
 
-    expect(history[1].name).toBe(`AFFiNE AI`);
+    expect(history[1].name).toBe(`AFFiNE AI\nwith your docs`);
-
+    expect(
-    // TODO(@akumatus): not stable
+      await page.locator('chat-panel affine-footnote-node').count()
-    // expect(history[1].name).toBe(`AFFiNE AI\nwith your docs`);
+    ).toBeGreaterThan(0);
-    // expect(
-    //   await page.locator('chat-panel affine-footnote-node').count()
-    // ).toBeGreaterThan(0);
     await clearChat(page);
     expect((await collectChat(page)).length).toBe(0);