feat(core): poll context docs and files embedding status (#10843)

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

### What Changed?
- Add status filed to `CopilotContextDoc` to querying document embedding processing progress.
- Change `ChipState` from `success` to `finished` to better align with backend server status.
- Add `pollContextDocsAndFiles` API for embedding status polling.

### About Polling
- Set the minimum interval to 1 second and the maximum interval to 30 seconds
- Use exponential backoff and increase the interval by 50% after each poll
- Make sure the interval does not exceed the maximum

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

@@ -242,6 +242,11 @@ declare global {
       ): AIActionTextResponse<T>;
     }
 
+    type AIDocsAndFilesContext = {
+      docs: CopilotContextDoc[];
+      files: CopilotContextFile[];
+    };
+
     interface AIContextService {
       createContext: (
         workspaceId: string,
@@ -274,13 +279,14 @@ declare global {
         workspaceId: string,
         sessionId: string,
         contextId: string
-      ) => Promise<
-        | {
-            docs: Array<CopilotContextDoc>;
-            files: Array<CopilotContextFile>;
-          }
-        | undefined
-      >;
+      ) => Promise<AIDocsAndFilesContext | undefined>;
+      pollContextDocsAndFiles: (
+        workspaceId: string,
+        sessionId: string,
+        contextId: string,
+        onPoll: (result: AIDocsAndFilesContext | undefined) => void,
+        abortSignal: AbortSignal
+      ) => Promise<void>;
       matchContext: (
         contextId: string,
         content: string,

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

@@ -70,13 +70,13 @@ export type ChatBlockMessage = ChatMessage & {
   avatarUrl?: string;
 };
 
-export type ChipState = 'candidate' | 'processing' | 'success' | 'failed';
+export type ChipState = 'candidate' | 'processing' | 'finished' | 'failed';
 
 export interface BaseChip {
   /**
    * candidate: the chip is a candidate for the chat
    * processing: the chip is processing
-   * success: the chip is successfully processed
+   * finished: the chip is successfully processed
    * failed: the chip is failed to process
    */
   state: ChipState;

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

@@ -68,6 +68,9 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<ChatContextValue>) => void;
 
+  @property({ attribute: false })
+  accessor pollContextDocsAndFiles!: () => void;
+
   @property({ attribute: false })
   accessor docDisplayConfig!: DocDisplayConfig;
 
@@ -201,6 +204,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
       await this._removeFromContext(chip);
     }
     await this._addToContext(chip);
+    this.pollContextDocsAndFiles();
   };
 
   private readonly _updateChip = (
@@ -229,19 +233,24 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
 
   private readonly _removeChip = async (chip: ChatChip) => {
     if (isDocChip(chip)) {
-      await this._removeFromContext(chip);
-      this.updateContext({
-        chips: this.chatContextValue.chips.filter(item => {
-          return !isDocChip(item) || item.docId !== chip.docId;
-        }),
-      });
-    } else {
-      await this._removeFromContext(chip);
-      this.updateContext({
-        chips: this.chatContextValue.chips.filter(item => {
-          return !isFileChip(item) || item.file !== chip.file;
-        }),
-      });
+      const removed = await this._removeFromContext(chip);
+      if (removed) {
+        this.updateContext({
+          chips: this.chatContextValue.chips.filter(item => {
+            return !isDocChip(item) || item.docId !== chip.docId;
+          }),
+        });
+      }
+    }
+    if (isFileChip(chip)) {
+      const removed = await this._removeFromContext(chip);
+      if (removed) {
+        this.updateContext({
+          chips: this.chatContextValue.chips.filter(item => {
+            return !isFileChip(item) || item.file !== chip.file;
+          }),
+        });
+      }
     }
   };
 
@@ -255,7 +264,8 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
         contextId,
         docId: chip.docId,
       });
-    } else {
+    }
+    if (isFileChip(chip)) {
       try {
         const blobId = await this.host.doc.blobSync.set(chip.file);
         const contextFile = await AIProvider.context.addContextFile(chip.file, {
@@ -263,7 +273,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
           blobId,
         });
         this._updateChip(chip, {
-          state: 'success',
+          state: contextFile.status,
           blobId: contextFile.blobId,
           fileId: contextFile.id,
         });
@@ -276,22 +286,26 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
     }
   };
 
-  private readonly _removeFromContext = async (chip: ChatChip) => {
+  private readonly _removeFromContext = async (
+    chip: ChatChip
+  ): Promise<boolean> => {
     const contextId = await this.getContextId();
     if (!contextId || !AIProvider.context) {
-      return;
+      return false;
     }
     if (isDocChip(chip)) {
-      await AIProvider.context.removeContextDoc({
+      return await AIProvider.context.removeContextDoc({
         contextId,
         docId: chip.docId,
       });
-    } else if (isFileChip(chip) && chip.fileId) {
-      await AIProvider.context.removeContextFile({
+    }
+    if (isFileChip(chip) && chip.fileId) {
+      return await AIProvider.context.removeContextFile({
         contextId,
         fileId: chip.fileId,
       });
     }
+    return true;
   };
 
   private readonly _checkTokenLimit = (
@@ -305,7 +319,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
       if (chip.docId === newChip.docId) {
         return acc + newTokenCount;
       }
-      if (chip.markdown?.value && chip.state === 'success') {
+      if (chip.markdown?.value && chip.state === 'finished') {
         const tokenCount =
           chip.tokenCount ?? estimateTokenCount(chip.markdown.value);
         return acc + tokenCount;

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

@@ -226,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 === 'success')
+      !!this.chatContextValue.chips.filter(chip => chip.state === 'finished')
         .length
     );
   }
@@ -561,7 +561,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
       : [];
     const contexts = this.chatContextValue.chips.reduce(
       (acc, chip, index) => {
-        if (chip.state !== 'success') {
+        if (chip.state !== 'finished') {
           return acc;
         }
         if (isDocChip(chip) && !!chip.markdown?.value) {

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

@@ -112,7 +112,7 @@ export class ChatPanelDocChip extends SignalWatcher(
         const markdown = this.chip.markdown ?? new Signal<string>('');
         markdown.value = value;
         this.updateChip(this.chip, {
-          state: 'success',
+          state: 'finished',
           markdown,
           tokenCount,
         });

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

@@ -1,6 +1,7 @@
 import './chat-panel-input';
 import './chat-panel-messages';
 
+import type { CopilotContextDoc, CopilotContextFile } from '@affine/graphql';
 import type { EditorHost } from '@blocksuite/affine/block-std';
 import { ShadowlessElement } from '@blocksuite/affine/block-std';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
@@ -35,7 +36,7 @@ import type {
   FileChip,
 } from './chat-context';
 import type { ChatPanelMessages } from './chat-panel-messages';
-import { isDocContext } from './components/utils';
+import { isDocChip, isDocContext } from './components/utils';
 
 const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
   quote: '',
@@ -174,7 +175,7 @@ export class ChatPanel extends SignalWatcher(
     this._scrollToEnd();
   };
 
-  private readonly _updateChips = async () => {
+  private readonly _initChips = async () => {
     // context not initialized, show candidate chip
     if (!this._chatSessionId || !this._chatContextId) {
       return;
@@ -191,12 +192,16 @@ export class ChatPanel extends SignalWatcher(
       (a, b) =>
         new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
     );
+    let allDone = true;
     const chips: ChatChip[] = await Promise.all(
       list.map(async item => {
+        if (item.status === 'processing') {
+          allDone = false;
+        }
         if (isDocContext(item)) {
           return {
             docId: item.id,
-            state: 'processing',
+            state: item.status || 'processing',
           } as DocChip;
         }
         const file = await this.host.doc.blobSync.get(item.blobId);
@@ -212,7 +217,7 @@ export class ChatPanel extends SignalWatcher(
             file: new File([file], item.name),
             blobId: item.blobId,
             fileId: item.id,
-            state: item.status === 'finished' ? 'success' : item.status,
+            state: item.status,
             tooltip: item.error,
           } as FileChip;
         }
@@ -223,6 +228,10 @@ export class ChatPanel extends SignalWatcher(
       ...this.chatContextValue,
       chips,
     };
+
+    if (!allDone) {
+      await this._pollContextDocsAndFiles();
+    }
   };
 
   private readonly _getSessionId = async () => {
@@ -285,6 +294,8 @@ export class ChatPanel extends SignalWatcher(
 
   private _width: Signal<number | undefined> = signal(undefined);
 
+  private _pollAbortController: AbortController | null = null;
+
   private readonly _scrollToEnd = () => {
     if (!this._wheelTriggered) {
       this._chatMessages.value?.scrollToEnd();
@@ -345,14 +356,79 @@ export class ChatPanel extends SignalWatcher(
           this._chatSessionId
         );
       }
-      await this._updateChips();
+      await this._initChips();
     } catch (error) {
       console.error(error);
     }
   };
 
+  private readonly _pollContextDocsAndFiles = async () => {
+    if (!this._chatSessionId || !this._chatContextId || !AIProvider.context) {
+      return;
+    }
+    if (this._pollAbortController) {
+      // already polling, return
+      return;
+    }
+    this._pollAbortController = new AbortController();
+    await AIProvider.context.pollContextDocsAndFiles(
+      this.doc.workspace.id,
+      this._chatSessionId,
+      this._chatContextId,
+      this._onPoll,
+      this._pollAbortController.signal
+    );
+  };
+
+  private readonly _onPoll = (
+    result?: BlockSuitePresets.AIDocsAndFilesContext
+  ) => {
+    if (!result) {
+      this._abortPoll();
+      return;
+    }
+    const { docs = [], files = [] } = result;
+    const hashMap = new Map<string, CopilotContextDoc | CopilotContextFile>();
+    let allDone = true;
+    docs.forEach(doc => {
+      hashMap.set(doc.id, doc);
+      if (doc.status === 'processing') {
+        allDone = false;
+      }
+    });
+    files.forEach(file => {
+      hashMap.set(file.id, file);
+      if (file.status === 'processing') {
+        allDone = false;
+      }
+    });
+    const nextChips = this.chatContextValue.chips.map(chip => {
+      const id = isDocChip(chip) ? chip.docId : chip.fileId;
+      const item = id && hashMap.get(id);
+      if (item && item.status) {
+        return {
+          ...chip,
+          state: item.status,
+        };
+      }
+      return chip;
+    });
+    this.updateContext({
+      chips: nextChips,
+    });
+    if (allDone) {
+      this._abortPoll();
+    }
+  };
+
+  private readonly _abortPoll = () => {
+    this._pollAbortController?.abort();
+    this._pollAbortController = null;
+  };
+
   protected override updated(_changedProperties: PropertyValues) {
     if (_changedProperties.has('doc')) {
+      this._abortPoll();
       this._chatSessionId = null;
       this._chatContextId = null;
       this.chatContextValue = DEFAULT_CHAT_CONTEXT_VALUE;
@@ -494,6 +570,7 @@ export class ChatPanel extends SignalWatcher(
         .chatContextValue=${this.chatContextValue}
         .getContextId=${this._getContextId}
         .updateContext=${this.updateContext}
+        .pollContextDocsAndFiles=${this._pollContextDocsAndFiles}
         .docDisplayConfig=${this.docDisplayConfig}
         .docSearchMenuConfig=${this.docSearchMenuConfig}
       ></chat-panel-chips>

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

@@ -451,6 +451,38 @@ Could you make a new website based on these notes and send back just the html fi
     ) => {
       return client.getContextDocsAndFiles(workspaceId, sessionId, contextId);
     },
+    pollContextDocsAndFiles: async (
+      workspaceId: string,
+      sessionId: string,
+      contextId: string,
+      onPoll: (
+        result: BlockSuitePresets.AIDocsAndFilesContext | undefined
+      ) => void,
+      abortSignal: AbortSignal
+    ) => {
+      const poll = async () => {
+        const result = await client.getContextDocsAndFiles(
+          workspaceId,
+          sessionId,
+          contextId
+        );
+        onPoll(result);
+      };
+
+      let attempts = 0;
+      const MIN_INTERVAL = 1000;
+      const MAX_INTERVAL = 30 * 1000;
+
+      while (!abortSignal.aborted) {
+        await poll();
+        const interval = Math.min(
+          MIN_INTERVAL * Math.pow(1.5, attempts),
+          MAX_INTERVAL
+        );
+        attempts++;
+        await new Promise(resolve => setTimeout(resolve, interval));
+      }
+    },
     matchContext: async (
       contextId: string,
       content: string,

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

@@ -1026,7 +1026,7 @@ test.describe('chat with doc', () => {
     expect(await chipTitle.textContent()).toBe('Untitled');
     let chip = await page.getByTestId('chat-panel-chip');
     // oxlint-disable-next-line unicorn/prefer-dom-node-dataset
-    expect(await chip.getAttribute('data-state')).toBe('success');
+    expect(await chip.getAttribute('data-state')).toBe('finished');
 
     const editorTitle = await page.locator('doc-title .inline-editor').nth(0);
     await editorTitle.pressSequentially('AFFiNE AI', {
@@ -1048,7 +1048,7 @@ test.describe('chat with doc', () => {
 
     expect(await chipTitle.textContent()).toBe('AFFiNE AI');
     // oxlint-disable-next-line unicorn/prefer-dom-node-dataset
-    expect(await chip.getAttribute('data-state')).toBe('success');
+    expect(await chip.getAttribute('data-state')).toBe('finished');
 
     await typeChatSequentially(page, 'What is AFiAI?');
     await page.keyboard.press('Enter');
@@ -1071,6 +1071,6 @@ test.describe('chat with doc', () => {
     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');
+    expect(await chip2.getAttribute('data-state')).toBe('finished');
   });
 });