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/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 => {
-        assertExists(AIProvider.actions.chat);
+      .then(async context => {
+        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 =
-  | 'candidate'
-  | 'uploading'
-  | 'embedding'
-  | 'success'
-  | 'failed';
+export type ChipState = 'candidate' | 'processing' | 'success' | 'failed';
 
 export interface BaseChip {
   /**
    * candidate: the chip is a candidate for the chat
-   * uploading: the chip is uploading
-   * embedding: the chip is embedding
-   * success: the chip is successfully embedded
-   * failed: the chip is failed to embed
+   * processing: the chip is processing
+   * success: the chip is successfully processed
+   * failed: the chip is failed to process
    */
   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)) {
-        return !isDocChip(chip) || item.docId !== chip.docId;
+      if (isDocChip(chip)) {
+        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) => {
-    this.updateContext({
-      chips: this.chatContextValue.chips.filter(item => {
-        if (isDocChip(item)) {
-          return !isDocChip(chip) || item.docId !== chip.docId;
-        } else {
-          return !isFileChip(chip) || item.fileId !== chip.fileId;
-        }
-      }),
-    });
+  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.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 =
-      !!this.chatContextValue.images.length ||
-      !!this.chatContextValue.chips.filter(chip => chip.state !== 'candidate')
-        .length;
-    const networkActive = !!this.networkSearchConfig.enabled.value;
-    const uploadDisabled = networkActive && !networkDisabled;
+    const uploadDisabled = this._isNetworkActive && !this._isNetworkDisabled;
     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}
-                  data-active=${networkActive}
-                  @click=${networkDisabled
+                  aria-disabled=${this._isNetworkDisabled}
+                  data-active=${this._isNetworkActive}
+                  @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 (!chatSessionId) return;
+      if (!this.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, {
-        state: 'embedding',
+      this.addChip({
+        ...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') {
-    return 'Embedding...';
-  }
-  if (state === 'uploading') {
-    return 'Uploading...';
+  if (state === 'processing') {
+    return 'Processing...';
   }
   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(() => {
-    const counter = ++this._resettingCounter;
+  private readonly _updateHistory = async () => {
+    const { doc } = this;
     this.isLoading = true;
-    (async () => {
-      const { doc } = this;
 
-      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 currentRequest = ++this._updateHistoryCounter;
 
-      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 history = histories.at(-1);
-        if (!history) return;
-        this.chatContextValue.chatSessionId = history.sessionId;
-        items.push(...history.messages);
-        AIProvider.LAST_ROOT_SESSION_ID = history.sessionId;
-      }
+    const items: ChatItem[] = actions ? [...actions] : [];
 
-      const { chips } = this.chatContextValue;
-      const defaultChip: DocChip = {
-        docId: this.doc.id,
-        state: 'candidate',
-      };
-      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,
-      };
+    if (histories?.at(-1)) {
+      const history = histories.at(-1);
+      if (!history) return;
+      items.push(...history.messages);
+      AIProvider.LAST_ROOT_SESSION_ID = history.sessionId;
+    }
 
-      this.isLoading = false;
-      this._scrollToEnd();
-    })().catch(console.error);
-  }, 200);
+    this.chatContextValue = {
+      ...this.chatContextValue,
+      items: items.sort(
+        (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(() => {
-        this.chatContextValue.chatSessionId = null;
-        // TODO get from CopilotContext
-        this.chatContextValue.chips = [];
-        this._resetItems();
+      requestAnimationFrame(async () => {
+        await this._initPanel();
       });
     }
 
@@ -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 => {
-        if (userInfo) {
-          this._resetItems();
-        }
+      AIProvider.slots.userInfo.on(async () => {
+        await this._initPanel();
       })
     );
     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;