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)
2026-02-15 05:37:32 +00:00 · 2025-02-12 08:33:06 +00:00
parent 53fdb1e8a5
commit 58fed5928b
21 changed files with 588 additions and 244 deletions
--- a/packages/frontend/core/src/blocksuite/presets/ai/actions/doc-handler.ts
+++ b/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,
--- a/packages/frontend/core/src/blocksuite/presets/ai/actions/types.ts
+++ b/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: (
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-context.ts
+++ b/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;
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-chips.ts
+++ b/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,
+      });
+    }
  };
 }
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts
+++ b/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 ?? '';
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-messages.ts
+++ b/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>`
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/add-popover.ts
+++ b/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();
  };
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/doc-chip.ts
+++ b/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);
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/file-chip.ts
+++ b/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);
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/components/utils.ts
+++ b/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;
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts
+++ b/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}
--- a/packages/frontend/core/src/blocksuite/presets/ai/provider.ts
+++ b/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;
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts
+++ b/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,
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
+++ b/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,
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx
+++ b/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,
-  networkSearchService: AINetworkSearchService
+  globalDialogService: GlobalDialogService
 ) {
-  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,
--- a/packages/frontend/core/src/components/providers/workspace-side-effects.tsx
+++ b/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,
-      networkSearchService
+      globalDialogService
    );
    return () => {
      dispose();
@@ -164,7 +161,6 @@ export const WorkspaceSideEffects = () => {
    fetchService,
    workspaceDialogService,
    graphqlService,
-    networkSearchService,
    globalDialogService,
  ]);