refactor(core): ai input (#11381)

Close [BS-2758](https://linear.app/affine-design/issue/BS-2758). Support [BS-2583](https://linear.app/affine-design/issue/BS-2583). ### What changed? - Extend `ChatPanelInput` and `ChatBlockInput` from the same abstract class `AIChatInput` to reduce duplication of code. - Unify the context interface of `chat-panel` and `chat-block`. - Rename `items` field to `messages`. - Remove duplicated type declare.
2026-02-12 04:18:54 +00:00 · 2025-04-02 13:37:30 +00:00
parent c8392d0cbe
commit c5c6978136
30 changed files with 777 additions and 1079 deletions
--- a/packages/frontend/core/src/blocksuite/ai/_common/chat-actions-handle.ts
+++ b/packages/frontend/core/src/blocksuite/ai/_common/chat-actions-handle.ts
@@ -41,7 +41,7 @@ import {
 import type { TemplateResult } from 'lit';

 import { insertFromMarkdown } from '../../utils';
-import type { ChatMessage } from '../blocks';
+import type { ChatMessage } from '../components/ai-chat-messages';
 import { AIProvider, type AIUserInfo } from '../provider';
 import { reportResponse } from '../utils/action-reporter';
 import { insertBelow } from '../utils/editor-actions';
--- a/packages/frontend/core/src/blocksuite/ai/blocks/ai-chat-block/ai-chat-block.ts
+++ b/packages/frontend/core/src/blocksuite/ai/blocks/ai-chat-block/ai-chat-block.ts
@@ -3,8 +3,9 @@ import { BlockComponent } from '@blocksuite/affine/std';
 import { computed } from '@preact/signals-core';
 import { html } from 'lit';

+import { ChatMessagesSchema } from '../../components/ai-chat-messages';
 import { ChatWithAIIcon } from './components/icon';
-import { type AIChatBlockModel, ChatMessagesSchema } from './model';
+import { type AIChatBlockModel } from './model';
 import { AIChatBlockStyles } from './styles';

@Peekable({
--- a/packages/frontend/core/src/blocksuite/ai/blocks/ai-chat-block/components/ai-chat-messages.ts
+++ b/packages/frontend/core/src/blocksuite/ai/blocks/ai-chat-block/components/ai-chat-messages.ts
@@ -6,7 +6,10 @@ import { property } from 'lit/decorators.js';
 import { classMap } from 'lit/directives/class-map.js';
 import { repeat } from 'lit/directives/repeat.js';

-import type { ChatMessage, MessageRole } from '../model';
+import type {
+  ChatMessage,
+  MessageRole,
+} from '../../../components/ai-chat-messages';
 import { UserInfoTemplate } from './user-info';

 export class AIChatMessage extends LitElement {
--- a/packages/frontend/core/src/blocksuite/ai/blocks/ai-chat-block/components/user-info.ts
+++ b/packages/frontend/core/src/blocksuite/ai/blocks/ai-chat-block/components/user-info.ts
@@ -2,7 +2,10 @@ import { baseTheme } from '@toeverything/theme';
 import { css, html, LitElement, type TemplateResult, unsafeCSS } from 'lit';
 import { property } from 'lit/decorators.js';

-import type { MessageRole, MessageUserInfo } from '../model';
+import type {
+  MessageRole,
+  MessageUserInfo,
+} from '../../../components/ai-chat-messages';
 import { AffineAIIcon } from './icon';

 export class UserInfo extends LitElement {
--- a/packages/frontend/core/src/blocksuite/ai/blocks/ai-chat-block/model/index.ts
+++ b/packages/frontend/core/src/blocksuite/ai/blocks/ai-chat-block/model/index.ts
@@ -1,3 +1,2 @@
 export * from './ai-chat-model';
 export * from './consts';
-export * from './types';
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/action-wrapper.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/action-wrapper.ts
@@ -23,8 +23,8 @@ import {
 import { css, html, LitElement, nothing, type TemplateResult } from 'lit';
 import { property, state } from 'lit/decorators.js';

+import { type ChatAction } from '../../components/ai-chat-messages';
 import { createTextRenderer } from '../../components/text-renderer';
-import type { ChatAction } from '../chat-context';
 import { HISTORY_IMAGE_ACTIONS } from '../const';

 const icons: Record<string, TemplateResult<1>> = {
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/image-to-text.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/image-to-text.ts
@@ -5,7 +5,7 @@ import { html, nothing } from 'lit';
 import { property } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';

-import type { ChatAction } from '../chat-context';
+import { type ChatAction } from '../../components/ai-chat-messages';

 export class ActionImageToText extends WithDisposable(ShadowlessElement) {
  @property({ attribute: false })
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/image.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/image.ts
@@ -8,7 +8,7 @@ import { html, nothing } from 'lit';
 import { property } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';

-import type { ChatAction } from '../chat-context';
+import { type ChatAction } from '../../components/ai-chat-messages';

 export class ActionImage extends WithDisposable(ShadowlessElement) {
  @property({ attribute: false })
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/make-real.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/make-real.ts
@@ -7,8 +7,8 @@ import { html } from 'lit';
 import { property } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';

+import { type ChatAction } from '../../components/ai-chat-messages';
 import { createIframeRenderer } from '../../messages/wrapper';
-import type { ChatAction } from '../chat-context';

 export class ActionMakeReal extends WithDisposable(ShadowlessElement) {
  @property({ attribute: false })
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/mindmap.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/mindmap.ts
@@ -7,7 +7,7 @@ import { html } from 'lit';
 import { property } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';

-import type { ChatAction } from '../chat-context';
+import { type ChatAction } from '../../components/ai-chat-messages';

 export class ActionMindmap extends WithDisposable(ShadowlessElement) {
  @property({ attribute: false })
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/slides.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/slides.ts
@@ -8,7 +8,7 @@ import { html, nothing } from 'lit';
 import { property } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';

-import type { ChatAction } from '../chat-context';
+import { type ChatAction } from '../../components/ai-chat-messages';

 export class ActionSlides extends WithDisposable(ShadowlessElement) {
  @property({ attribute: false })
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/text.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/text.ts
@@ -7,8 +7,9 @@ import { css, html, LitElement } from 'lit';
 import { property } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';

+import { type ChatAction } from '../../components/ai-chat-messages';
 import { createTextRenderer } from '../../components/text-renderer';
-import type { ChatAction } from '../chat-context';
+
 export class ActionText extends WithDisposable(LitElement) {
  static override styles = css`
    .original-text {
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-config.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-config.ts
@@ -10,9 +10,3 @@ export interface AppSidebarConfig {
    cleanup: () => void;
  };
 }
-
-export interface AINetworkSearchConfig {
-  visible: Signal<boolean | undefined>;
-  enabled: Signal<boolean | undefined>;
-  setEnabled: (state: boolean) => void;
-}
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts
@@ -1,40 +1,12 @@
+import type {
+  ChatStatus,
+  HistoryMessage,
+} from '../components/ai-chat-messages';
 import type { AIError } from '../provider';

-export type ChatMessage = {
-  id: string;
-  content: string;
-  role: 'user' | 'assistant';
-  attachments?: string[];
-  createdAt: string;
-};
-
-export type ChatAction = {
-  action: string;
-  messages: ChatMessage[];
-  sessionId: string;
-  createdAt: string;
-};
-
-export type ChatItem = ChatMessage | ChatAction;
-
-export function isChatAction(item: ChatItem): item is ChatAction {
-  return 'action' in item;
-}
-
-export function isChatMessage(item: ChatItem): item is ChatMessage {
-  return 'role' in item;
-}
-
-export type ChatStatus =
-  | 'loading'
-  | 'success'
-  | 'error'
-  | 'idle'
-  | 'transmitting';
-
 export type ChatContextValue = {
  // history messages of the chat
-  items: ChatItem[];
+  messages: HistoryMessage[];
  status: ChatStatus;
  error: AIError | null;
  // plain-text of the selected content
@@ -45,9 +17,3 @@ export type ChatContextValue = {
  images: File[];
  abortController: AbortController | null;
 };
-
-export type ChatBlockMessage = ChatMessage & {
-  userId?: string;
-  userName?: string;
-  avatarUrl?: string;
-};
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-input.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-input.ts
@@ -1,476 +1,9 @@
-import { stopPropagation } from '@affine/core/utils';
-import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
-import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
-import { openFileOrFiles } from '@blocksuite/affine/shared/utils';
-import type { EditorHost } from '@blocksuite/affine/std';
-import {
-  BroomIcon,
-  CloseIcon,
-  ImageIcon,
-  PublishIcon,
-} from '@blocksuite/icons/lit';
-import { css, html, LitElement, nothing } from 'lit';
-import { property, query, state } from 'lit/decorators.js';
-import { repeat } from 'lit/directives/repeat.js';
-
-import { ChatAbortIcon, ChatSendIcon } from '../_common/icons';
-import type {
-  ChatChip,
-  DocDisplayConfig,
-  FileChip,
-} from '../components/ai-chat-chips';
-import { isDocChip, isFileChip } from '../components/ai-chat-chips';
+import { AIChatInput } from '../components/ai-chat-input';
+import type { ChatMessage } from '../components/ai-chat-messages';
 import { type AIError, AIProvider } from '../provider';
-import { reportResponse } from '../utils/action-reporter';
 import { readBlobAsURL } from '../utils/image';
-import type { AINetworkSearchConfig } from './chat-config';
-import type { ChatContextValue, ChatMessage } from './chat-context';
-import { PROMPT_NAME_AFFINE_AI, PROMPT_NAME_NETWORK_SEARCH } from './const';
-
-const MaximumImageCount = 32;
-
-function getFirstTwoLines(text: string) {
-  const lines = text.split('\n');
-  return lines.slice(0, 2);
-}
-
-export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
-  static override styles = css`
-    .chat-panel-input {
-      display: flex;
-      flex-direction: column;
-      justify-content: space-between;
-      gap: 12px;
-      position: relative;
-      margin-top: 12px;
-      border-radius: 4px;
-      padding: 8px;
-      min-height: 94px;
-      box-sizing: border-box;
-      border-width: 1px;
-      border-style: solid;
-
-      .chat-selection-quote {
-        padding: 4px 0px 8px 0px;
-        padding-left: 15px;
-        max-height: 56px;
-        font-size: 14px;
-        font-weight: 400;
-        line-height: 22px;
-        color: var(--affine-text-secondary-color);
-        position: relative;
-
-        div {
-          white-space: nowrap;
-          overflow: hidden;
-          text-overflow: ellipsis;
-        }
-
-        .chat-quote-close {
-          position: absolute;
-          right: 0;
-          top: 0;
-          cursor: pointer;
-          display: none;
-          width: 16px;
-          height: 16px;
-          border-radius: 4px;
-          border: 1px solid var(--affine-border-color);
-          background-color: var(--affine-white);
-        }
-      }
-
-      .chat-selection-quote:hover .chat-quote-close {
-        display: flex;
-        justify-content: center;
-        align-items: center;
-      }
-
-      .chat-selection-quote::after {
-        content: '';
-        width: 2px;
-        height: calc(100% - 10px);
-        margin-top: 5px;
-        position: absolute;
-        left: 0;
-        top: 0;
-        background: var(--affine-quote-color);
-        border-radius: 18px;
-      }
-    }
-
-    .chat-panel-input-actions {
-      display: flex;
-      gap: 8px;
-      align-items: center;
-
-      div {
-        width: 24px;
-        height: 24px;
-        cursor: pointer;
-      }
-
-      div:nth-child(2) {
-        margin-left: auto;
-      }
-
-      .image-upload,
-      .chat-history-clear,
-      .chat-network-search {
-        display: flex;
-        justify-content: center;
-        align-items: center;
-        svg {
-          width: 20px;
-          height: 20px;
-          color: ${unsafeCSSVarV2('icon/primary')};
-        }
-      }
-      .chat-history-clear svg {
-        color: var(--affine-text-secondary-color);
-      }
-      .chat-network-search[data-active='true'] svg {
-        color: ${unsafeCSSVarV2('icon/activated')};
-      }
-
-      .image-upload[aria-disabled='true'],
-      .chat-network-search[aria-disabled='true'] {
-        cursor: not-allowed;
-      }
-      .image-upload[aria-disabled='true'] svg,
-      .chat-network-search[aria-disabled='true'] svg {
-        color: var(--affine-text-disable-color) !important;
-      }
-    }
-
-    .chat-panel-input {
-      textarea {
-        width: 100%;
-        padding: 0;
-        margin: 0;
-        border: none;
-        line-height: 22px;
-        font-size: var(--affine-font-sm);
-        font-weight: 400;
-        font-family: var(--affine-font-family);
-        color: var(--affine-text-primary-color);
-        box-sizing: border-box;
-        resize: none;
-        overflow-y: hidden;
-      }
-
-      textarea::placeholder {
-        font-size: 14px;
-        font-weight: 400;
-        font-family: var(--affine-font-family);
-        color: var(--affine-placeholder-color);
-      }
-
-      textarea:focus {
-        outline: none;
-      }
-    }
-
-    .chat-panel-send svg rect {
-      fill: var(--affine-primary-color);
-    }
-    .chat-panel-send[aria-disabled='true'] {
-      cursor: not-allowed;
-    }
-    .chat-panel-send[aria-disabled='true'] svg rect {
-      fill: var(--affine-text-disable-color);
-    }
-  `;
-
-  @property({ attribute: false })
-  accessor host!: EditorHost;
-
-  @query('image-preview-grid')
-  accessor imagePreviewGrid: HTMLDivElement | null = null;
-
-  @query('textarea')
-  accessor textarea!: HTMLTextAreaElement;
-
-  @state()
-  accessor isInputEmpty = true;
-
-  @state()
-  accessor focused = false;
-
-  @property({ attribute: false })
-  accessor chatContextValue!: ChatContextValue;
-
-  @property({ attribute: false })
-  accessor chips: ChatChip[] = [];
-
-  @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;
-
-  @property({ attribute: false })
-  accessor cleanupHistories!: () => Promise<void>;
-
-  @property({ attribute: false })
-  accessor networkSearchConfig!: AINetworkSearchConfig;
-
-  @property({ attribute: false })
-  accessor docDisplayConfig!: DocDisplayConfig;
-
-  @property({ attribute: 'data-testid', reflect: true })
-  accessor testId = 'chat-panel-input-container';
-
-  private get _isNetworkActive() {
-    return (
-      !!this.networkSearchConfig.visible.value &&
-      !!this.networkSearchConfig.enabled.value
-    );
-  }
-
-  private get _isNetworkDisabled() {
-    return (
-      !!this.chatContextValue.images.length ||
-      !!this.chips.filter(chip => chip.state === 'finished').length
-    );
-  }
-
-  private _getPromptName() {
-    if (this._isNetworkDisabled) {
-      return PROMPT_NAME_AFFINE_AI;
-    }
-    return this._isNetworkActive
-      ? PROMPT_NAME_NETWORK_SEARCH
-      : PROMPT_NAME_AFFINE_AI;
-  }
-
-  private async _updatePromptName(promptName: string) {
-    const sessionId = await this.getSessionId();
-    if (sessionId && AIProvider.session) {
-      await AIProvider.session.updateSession(sessionId, promptName);
-    }
-  }
-
-  private _addImages(images: File[]) {
-    const oldImages = this.chatContextValue.images;
-    this.updateContext({
-      images: [...oldImages, ...images].slice(0, MaximumImageCount),
-    });
-  }
-
-  private readonly _handleImageRemove = (index: number) => {
-    const oldImages = this.chatContextValue.images;
-    const newImages = oldImages.filter((_, i) => i !== index);
-    this.updateContext({ images: newImages });
-  };
-
-  private readonly _toggleNetworkSearch = (e: MouseEvent) => {
-    e.preventDefault();
-    e.stopPropagation();
-
-    const enable = this.networkSearchConfig.enabled.value;
-    this.networkSearchConfig.setEnabled(!enable);
-  };
-
-  private readonly _uploadImageFiles = async (_e: MouseEvent) => {
-    const images = await openFileOrFiles({
-      acceptType: 'Images',
-      multiple: true,
-    });
-    if (!images) return;
-    this._addImages(images);
-  };
-
-  override connectedCallback() {
-    super.connectedCallback();
-
-    this._disposables.add(
-      AIProvider.slots.requestSendWithChat.subscribe(
-        ({ input, context, host }) => {
-          if (this.host === host) {
-            context && this.updateContext(context);
-            const { updateComplete, send } = this;
-            updateComplete
-              .then(() => {
-                return send(input);
-              })
-              .catch(console.error);
-          }
-        }
-      )
-    );
-  }
-
-  protected override render() {
-    const { images, status } = this.chatContextValue;
-    const hasImages = images.length > 0;
-    const maxHeight = hasImages ? 272 + 2 : 200 + 2;
-    const uploadDisabled = this._isNetworkActive && !this._isNetworkDisabled;
-    return html`<style>
-        .chat-panel-input {
-          border-color: ${this.focused
-            ? 'var(--affine-primary-color)'
-            : 'var(--affine-border-color)'};
-          box-shadow: ${this.focused ? 'var(--affine-active-shadow)' : 'none'};
-          max-height: ${maxHeight}px !important;
-          user-select: none;
-        }
-      </style>
-      <div
-        class="chat-panel-input"
-        @pointerdown=${(e: MouseEvent) => {
-          if (e.target !== this.textarea) {
-            // by default the div will be focused and will blur the textarea
-            e.preventDefault();
-            this.textarea.focus();
-          }
-        }}
-      >
-        ${hasImages
-          ? html`
-              <image-preview-grid
-                .images=${images}
-                .onImageRemove=${this._handleImageRemove}
-              ></image-preview-grid>
-            `
-          : nothing}
-        ${this.chatContextValue.quote
-          ? html`<div
-              class="chat-selection-quote"
-              data-testid="chat-selection-quote"
-            >
-              ${repeat(
-                getFirstTwoLines(this.chatContextValue.quote),
-                line => line,
-                line => html`<div>${line}</div>`
-              )}
-              <div
-                class="chat-quote-close"
-                @click=${() => {
-                  this.updateContext({ quote: '', markdown: '' });
-                }}
-              >
-                ${CloseIcon()}
-              </div>
-            </div>`
-          : nothing}
-        <textarea
-          rows="1"
-          placeholder="What are your thoughts?"
-          @input=${() => {
-            const { textarea } = this;
-            this.isInputEmpty = !textarea.value.trim();
-            textarea.style.height = 'auto';
-            textarea.style.height = textarea.scrollHeight + 'px';
-            let imagesHeight = this.imagePreviewGrid?.scrollHeight ?? 0;
-            if (imagesHeight) imagesHeight += 12;
-            if (this.scrollHeight >= 200 + imagesHeight) {
-              textarea.style.height = '148px';
-              textarea.style.overflowY = 'scroll';
-            }
-          }}
-          @keydown=${async (evt: KeyboardEvent) => {
-            if (evt.key === 'Enter' && !evt.shiftKey && !evt.isComposing) {
-              await this._onTextareaSend(evt);
-            }
-          }}
-          @focus=${() => {
-            this.focused = true;
-          }}
-          @blur=${() => {
-            this.focused = false;
-          }}
-          @paste=${(event: ClipboardEvent) => {
-            const items = event.clipboardData?.items;
-            if (!items) return;
-
-            for (const index in items) {
-              const item = items[index];
-              if (item.kind === 'file' && item.type.indexOf('image') >= 0) {
-                const blob = item.getAsFile();
-                if (!blob) continue;
-                this._addImages([blob]);
-              }
-            }
-          }}
-          data-testid="chat-panel-input"
-        ></textarea>
-        <div class="chat-panel-input-actions">
-          <div
-            class="chat-history-clear"
-            @click=${async () => {
-              await this.cleanupHistories();
-            }}
-            data-testid="chat-panel-clear"
-          >
-            ${BroomIcon()}
-          </div>
-          ${this.networkSearchConfig.visible.value
-            ? html`
-                <div
-                  class="chat-network-search"
-                  data-testid="chat-network-search"
-                  aria-disabled=${this._isNetworkDisabled}
-                  data-active=${this._isNetworkActive}
-                  @click=${this._isNetworkDisabled
-                    ? undefined
-                    : this._toggleNetworkSearch}
-                  @pointerdown=${stopPropagation}
-                >
-                  ${PublishIcon()}
-                </div>
-              `
-            : nothing}
-          ${images.length < MaximumImageCount
-            ? html`<div
-                data-testid="chat-panel-input-image-upload"
-                class="image-upload"
-                aria-disabled=${uploadDisabled}
-                @click=${uploadDisabled ? undefined : this._uploadImageFiles}
-              >
-                ${ImageIcon()}
-              </div>`
-            : nothing}
-          ${status === 'transmitting'
-            ? html`<div
-                @click=${() => {
-                  this.chatContextValue.abortController?.abort();
-                  this.updateContext({ status: 'success' });
-                  reportResponse('aborted:stop');
-                }}
-                data-testid="chat-panel-stop"
-              >
-                ${ChatAbortIcon}
-              </div>`
-            : html`<div
-                @click="${this._onTextareaSend}"
-                class="chat-panel-send"
-                aria-disabled=${this.isInputEmpty}
-                data-testid="chat-panel-send"
-              >
-                ${ChatSendIcon}
-              </div>`}
-        </div>
-      </div>`;
-  }
-
-  private readonly _onTextareaSend = async (e: MouseEvent | KeyboardEvent) => {
-    e.preventDefault();
-    e.stopPropagation();
-
-    const value = this.textarea.value.trim();
-    if (value.length === 0) return;
-
-    this.textarea.value = '';
-    this.isInputEmpty = true;
-    this.textarea.style.height = 'unset';
-
-    await this.send(value);
-  };

+export class ChatPanelInput extends AIChatInput {
  send = async (text: string) => {
    const { status, markdown, images } = this.chatContextValue;
    if (status === 'loading' || status === 'transmitting') return;
@@ -478,7 +11,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {

    try {
      const { doc } = this.host;
-      const promptName = this._getPromptName();
+      const promptName = this.getPromptName();

      this.updateContext({
        images: [],
@@ -494,8 +27,8 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {

      const userInput = (markdown ? `${markdown}\n` : '') + text;
      this.updateContext({
-        items: [
-          ...this.chatContextValue.items,
+        messages: [
+          ...this.chatContextValue.messages,
          {
            id: '',
            role: 'user',
@@ -514,12 +47,13 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {

      // must update prompt name after local chat message is updated
      // otherwise, the unauthorized error can not be rendered properly
-      await this._updatePromptName(promptName);
+      await this.updatePromptName(promptName);

      const abortController = new AbortController();
      const sessionId = await this.getSessionId();
+      if (!sessionId) return;

-      const contexts = await this._getMatchedContexts(userInput);
+      const contexts = await this.getMatchedContexts(userInput);
      const stream = AIProvider.actions.chat?.({
        sessionId,
        input: userInput,
@@ -539,16 +73,16 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
        this.updateContext({ abortController });

        for await (const text of stream) {
-          const items = [...this.chatContextValue.items];
-          const last = items[items.length - 1] as ChatMessage;
+          const messages = [...this.chatContextValue.messages];
+          const last = messages[messages.length - 1] as ChatMessage;
          last.content += text;
-          this.updateContext({ items, status: 'transmitting' });
+          this.updateContext({ messages, status: 'transmitting' });
        }

        this.updateContext({ status: 'success' });

-        const { items } = this.chatContextValue;
-        const last = items[items.length - 1] as ChatMessage;
+        const { messages } = this.chatContextValue;
+        const last = messages[messages.length - 1] as ChatMessage;
        if (!last.id) {
          const historyIds = await AIProvider.histories?.ids(
            doc.workspace.id,
@@ -565,89 +99,6 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
      this.updateContext({ abortController: null });
    }
  };
-
-  private async _getMatchedContexts(userInput: string) {
-    const contextId = await this.getContextId();
-    if (!contextId) {
-      return { files: [], docs: [] };
-    }
-
-    const docContexts = new Map<
-      string,
-      { docId: string; docContent: string }
-    >();
-    const fileContexts = new Map<
-      string,
-      BlockSuitePresets.AIFileContextOption
-    >();
-
-    const { files: matchedFiles = [], docs: matchedDocs = [] } =
-      (await AIProvider.context?.matchContext(contextId, userInput)) ?? {};
-
-    matchedDocs.forEach(doc => {
-      docContexts.set(doc.docId, {
-        docId: doc.docId,
-        docContent: doc.content,
-      });
-    });
-
-    matchedFiles.forEach(file => {
-      const context = fileContexts.get(file.fileId);
-      if (context) {
-        context.fileContent += `\n${file.content}`;
-      } else {
-        const fileChip = this.chips.find(
-          chip => isFileChip(chip) && chip.fileId === file.fileId
-        ) as FileChip | undefined;
-        if (fileChip && fileChip.blobId) {
-          fileContexts.set(file.fileId, {
-            blobId: fileChip.blobId,
-            fileName: fileChip.file.name,
-            fileType: fileChip.file.type,
-            fileContent: file.content,
-          });
-        }
-      }
-    });
-
-    this.chips.forEach(chip => {
-      if (isDocChip(chip) && !!chip.markdown?.value) {
-        docContexts.set(chip.docId, {
-          docId: chip.docId,
-          docContent: chip.markdown.value,
-        });
-      }
-    });
-
-    const docs: BlockSuitePresets.AIDocContextOption[] = Array.from(
-      docContexts.values()
-    ).map(doc => {
-      const docMeta = this.docDisplayConfig.getDocMeta(doc.docId);
-      const docTitle = this.docDisplayConfig.getTitle(doc.docId);
-      const tags = docMeta?.tags
-        ? docMeta.tags
-            .map(tagId => this.docDisplayConfig.getTagTitle(tagId))
-            .join(',')
-        : '';
-      return {
-        docId: doc.docId,
-        docContent: doc.docContent,
-        docTitle,
-        tags,
-        createDate: docMeta?.createDate
-          ? new Date(docMeta.createDate).toISOString()
-          : '',
-        updatedDate: docMeta?.updatedDate
-          ? new Date(docMeta.updatedDate).toISOString()
-          : '',
-      };
-    });
-
-    return {
-      docs,
-      files: Array.from(fileContexts.values()),
-    };
-  }
 }

 declare global {
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-messages.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-messages.ts
@@ -14,13 +14,13 @@ import { repeat } from 'lit/directives/repeat.js';
 import { debounce } from 'lodash-es';

 import { AffineIcon } from '../_common/icons';
-import { type AIError, AIProvider, UnauthorizedError } from '../provider';
 import {
-  type ChatContextValue,
  type ChatMessage,
  isChatAction,
  isChatMessage,
-} from './chat-context';
+} from '../components/ai-chat-messages';
+import { type AIError, AIProvider, UnauthorizedError } from '../provider';
+import { type ChatContextValue } from './chat-context';
 import { HISTORY_IMAGE_ACTIONS } from './const';
 import { AIPreloadConfig } from './preload-config';

@@ -209,9 +209,9 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
  };

  protected override render() {
-    const { items, status, error } = this.chatContextValue;
+    const { messages, status, error } = this.chatContextValue;
    const { isLoading } = this;
-    const filteredItems = items.filter(item => {
+    const filteredItems = messages.filter(item => {
      return (
        isChatMessage(item) ||
        item.messages?.length === 3 ||
@@ -351,13 +351,13 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
      if (!sessionId) return;

      const abortController = new AbortController();
-      const items = [...this.chatContextValue.items];
-      const last = items[items.length - 1];
+      const messages = [...this.chatContextValue.messages];
+      const last = messages[messages.length - 1];
      if ('content' in last) {
        last.content = '';
        last.createdAt = new Date().toISOString();
      }
-      this.updateContext({ items, status: 'loading', error: null });
+      this.updateContext({ messages, status: 'loading', error: null });

      const stream = AIProvider.actions.chat?.({
        sessionId,
@@ -375,10 +375,10 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
      if (stream) {
        this.updateContext({ abortController });
        for await (const text of stream) {
-          const items = [...this.chatContextValue.items];
-          const last = items[items.length - 1] as ChatMessage;
+          const messages = [...this.chatContextValue.messages];
+          const last = messages[messages.length - 1] as ChatMessage;
          last.content += text;
-          this.updateContext({ items, status: 'transmitting' });
+          this.updateContext({ messages, status: 'transmitting' });
        }

        this.updateContext({ status: 'success' });
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/const.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/const.ts
@@ -8,6 +8,3 @@ export const HISTORY_IMAGE_ACTIONS = [
  'Remove background',
  'Convert to sticker',
 ];
-
-export const PROMPT_NAME_AFFINE_AI = 'Chat With AFFiNE AI';
-export const PROMPT_NAME_NETWORK_SEARCH = 'Search With AFFiNE AI';
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
@@ -35,21 +35,23 @@ import {
  isDocChip,
  isTagChip,
 } from '../components/ai-chat-chips';
+import type { AINetworkSearchConfig } from '../components/ai-chat-input';
+import { type HistoryMessage } from '../components/ai-chat-messages';
 import { AIProvider } from '../provider';
 import { extractSelectedContent } from '../utils/extract';
 import {
  getSelectedImagesAsBlobs,
  getSelectedTextContent,
 } from '../utils/selection-utils';
-import type { AINetworkSearchConfig, AppSidebarConfig } from './chat-config';
-import type { ChatContextValue, ChatItem } from './chat-context';
+import type { AppSidebarConfig } from './chat-config';
+import type { ChatContextValue } from './chat-context';
 import type { ChatPanelMessages } from './chat-panel-messages';

 const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
  quote: '',
  images: [],
  abortController: null,
-  items: [],
+  messages: [],
  status: 'idle',
  error: null,
  markdown: '',
@@ -160,19 +162,19 @@ export class ChatPanel extends SignalWatcher(
      return;
    }

-    const items: ChatItem[] = actions ? [...actions] : [];
+    const messages: HistoryMessage[] = actions ? [...actions] : [];

    const history = histories?.find(
      history => history.sessionId === this._chatSessionId
    );
    if (history) {
-      items.push(...history.messages);
+      messages.push(...history.messages);
      AIProvider.LAST_ROOT_SESSION_ID = history.sessionId;
    }

    this.chatContextValue = {
      ...this.chatContextValue,
-      items: items.sort(
+      messages: messages.sort(
        (a, b) =>
          new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
      ),
@@ -339,7 +341,7 @@ export class ChatPanel extends SignalWatcher(
          cancelText: 'Cancel',
        })
      ) {
-        const actionIds = this.chatContextValue.items
+        const actionIds = this.chatContextValue.messages
          .filter(item => 'sessionId' in item)
          .map(item => item.sessionId);
        await AIProvider.histories?.cleanup(
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/message/action.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/message/action.ts
@@ -6,7 +6,7 @@ import { ShadowlessElement } from '@blocksuite/affine/std';
 import { html } from 'lit';
 import { property } from 'lit/decorators.js';

-import { type ChatAction } from '../chat-context';
+import { type ChatAction } from '../../components/ai-chat-messages';
 import { HISTORY_IMAGE_ACTIONS } from '../const';

 export class ChatMessageAction extends WithDisposable(ShadowlessElement) {
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/message/assistant.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/message/assistant.ts
@@ -12,9 +12,12 @@ import {
  EdgelessEditorActions,
  PageEditorActions,
 } from '../../_common/chat-actions-handle';
+import {
+  type ChatMessage,
+  isChatMessage,
+} from '../../components/ai-chat-messages';
 import { AIChatErrorRenderer } from '../../messages/error';
 import { type AIError } from '../../provider';
-import { type ChatMessage, isChatMessage } from '../chat-context';

 export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
  static override styles = css`
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/message/user.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/message/user.ts
@@ -6,7 +6,7 @@ import { ShadowlessElement } from '@blocksuite/affine/std';
 import { css, html, nothing } from 'lit';
 import { property } from 'lit/decorators.js';

-import { type ChatMessage } from '../chat-context';
+import { type ChatMessage } from '../../components/ai-chat-messages';

 export class ChatMessageUser extends WithDisposable(ShadowlessElement) {
  static override styles = css`
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/ai-chat-input.ts
@@ -0,0 +1,588 @@
+import { stopPropagation } from '@affine/core/utils';
+import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
+import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
+import { openFileOrFiles } from '@blocksuite/affine/shared/utils';
+import type { EditorHost } from '@blocksuite/affine/std';
+import {
+  BroomIcon,
+  CloseIcon,
+  ImageIcon,
+  PublishIcon,
+} from '@blocksuite/icons/lit';
+import { css, html, LitElement, nothing } from 'lit';
+import { property, query, state } from 'lit/decorators.js';
+import { repeat } from 'lit/directives/repeat.js';
+import { styleMap } from 'lit/directives/style-map.js';
+
+import { ChatAbortIcon, ChatSendIcon } from '../../_common/icons';
+import { AIProvider } from '../../provider';
+import { reportResponse } from '../../utils/action-reporter';
+import type {
+  ChatChip,
+  DocDisplayConfig,
+  FileChip,
+} from '../ai-chat-chips/type';
+import { isDocChip, isFileChip } from '../ai-chat-chips/utils';
+import { PROMPT_NAME_AFFINE_AI, PROMPT_NAME_NETWORK_SEARCH } from './const';
+import type { AIChatInputContext, AINetworkSearchConfig } from './type';
+
+const MaximumImageCount = 32;
+
+function getFirstTwoLines(text: string) {
+  const lines = text.split('\n');
+  return lines.slice(0, 2);
+}
+
+export abstract class AIChatInput extends SignalWatcher(
+  WithDisposable(LitElement)
+) {
+  static override styles = css`
+    :host {
+      width: 100%;
+    }
+    .chat-panel-input {
+      display: flex;
+      flex-direction: column;
+      justify-content: space-between;
+      gap: 12px;
+      position: relative;
+      margin-top: 12px;
+      border-radius: 4px;
+      padding: 8px;
+      min-height: 94px;
+      box-sizing: border-box;
+      border-width: 1px;
+      border-style: solid;
+      border-color: var(--affine-border-color);
+
+      .chat-selection-quote {
+        padding: 4px 0px 8px 0px;
+        padding-left: 15px;
+        max-height: 56px;
+        font-size: 14px;
+        font-weight: 400;
+        line-height: 22px;
+        color: var(--affine-text-secondary-color);
+        position: relative;
+
+        div {
+          white-space: nowrap;
+          overflow: hidden;
+          text-overflow: ellipsis;
+        }
+
+        .chat-quote-close {
+          position: absolute;
+          right: 0;
+          top: 0;
+          cursor: pointer;
+          display: none;
+          width: 16px;
+          height: 16px;
+          border-radius: 4px;
+          border: 1px solid var(--affine-border-color);
+          background-color: var(--affine-white);
+        }
+      }
+
+      .chat-selection-quote:hover .chat-quote-close {
+        display: flex;
+        justify-content: center;
+        align-items: center;
+      }
+
+      .chat-selection-quote::after {
+        content: '';
+        width: 2px;
+        height: calc(100% - 10px);
+        margin-top: 5px;
+        position: absolute;
+        left: 0;
+        top: 0;
+        background: var(--affine-quote-color);
+        border-radius: 18px;
+      }
+    }
+
+    .chat-panel-input-actions {
+      display: flex;
+      gap: 8px;
+      align-items: center;
+
+      div {
+        width: 24px;
+        height: 24px;
+        cursor: pointer;
+      }
+
+      div:nth-child(2) {
+        margin-left: auto;
+      }
+
+      .image-upload,
+      .chat-history-clear,
+      .chat-network-search {
+        display: flex;
+        justify-content: center;
+        align-items: center;
+        svg {
+          width: 20px;
+          height: 20px;
+          color: ${unsafeCSSVarV2('icon/primary')};
+        }
+      }
+
+      .chat-history-clear svg {
+        color: var(--affine-text-secondary-color);
+      }
+      .chat-network-search[data-active='true'] svg {
+        color: ${unsafeCSSVarV2('icon/activated')};
+      }
+
+      .chat-history-clear[aria-disabled='true'],
+      .image-upload[aria-disabled='true'],
+      .chat-network-search[aria-disabled='true'] {
+        cursor: not-allowed;
+      }
+
+      .chat-history-clear[aria-disabled='true'] svg,
+      .image-upload[aria-disabled='true'] svg,
+      .chat-network-search[aria-disabled='true'] svg {
+        color: var(--affine-text-disable-color) !important;
+      }
+    }
+
+    .chat-panel-input {
+      textarea {
+        width: 100%;
+        padding: 0;
+        margin: 0;
+        border: none;
+        line-height: 22px;
+        font-size: var(--affine-font-sm);
+        font-weight: 400;
+        font-family: var(--affine-font-family);
+        color: var(--affine-text-primary-color);
+        box-sizing: border-box;
+        resize: none;
+        overflow-y: hidden;
+        background-color: transparent;
+      }
+
+      textarea::placeholder {
+        font-size: 14px;
+        font-weight: 400;
+        font-family: var(--affine-font-family);
+        color: var(--affine-placeholder-color);
+      }
+
+      textarea:focus {
+        outline: none;
+      }
+    }
+
+    .chat-panel-input[data-if-focused='true'] {
+      border-color: var(--affine-primary-color);
+      box-shadow: var(--affine-active-shadow);
+      user-select: none;
+    }
+
+    .chat-panel-send svg rect {
+      fill: var(--affine-primary-color);
+    }
+    .chat-panel-send[aria-disabled='true'] {
+      cursor: not-allowed;
+    }
+    .chat-panel-send[aria-disabled='true'] svg rect {
+      fill: var(--affine-text-disable-color);
+    }
+  `;
+
+  @property({ attribute: false })
+  accessor host!: EditorHost;
+
+  @query('image-preview-grid')
+  accessor imagePreviewGrid: HTMLDivElement | null = null;
+
+  @query('textarea')
+  accessor textarea!: HTMLTextAreaElement;
+
+  @state()
+  accessor isInputEmpty = true;
+
+  @state()
+  accessor focused = false;
+
+  @property({ attribute: false })
+  accessor chatContextValue!: AIChatInputContext;
+
+  @property({ attribute: false })
+  accessor chips: ChatChip[] = [];
+
+  @property({ attribute: false })
+  accessor getSessionId!: () => Promise<string | null | undefined>;
+
+  @property({ attribute: false })
+  accessor getContextId!: () => Promise<string | null | undefined>;
+
+  @property({ attribute: false })
+  accessor updateContext!: (context: Partial<AIChatInputContext>) => void;
+
+  @property({ attribute: false })
+  accessor cleanupHistories!: () => Promise<void>;
+
+  @property({ attribute: false })
+  accessor networkSearchConfig!: AINetworkSearchConfig;
+
+  @property({ attribute: false })
+  accessor docDisplayConfig!: DocDisplayConfig;
+
+  @property({ attribute: 'data-testid', reflect: true })
+  accessor testId = 'chat-panel-input-container';
+
+  private get _isNetworkActive() {
+    return (
+      !!this.networkSearchConfig.visible.value &&
+      !!this.networkSearchConfig.enabled.value
+    );
+  }
+
+  private get _isNetworkDisabled() {
+    return (
+      !!this.chatContextValue.images.length ||
+      !!this.chips.filter(chip => chip.state === 'finished').length
+    );
+  }
+
+  private get _isClearDisabled() {
+    return (
+      this.chatContextValue.status === 'loading' ||
+      this.chatContextValue.status === 'transmitting' ||
+      !this.chatContextValue.messages.length
+    );
+  }
+
+  protected getPromptName() {
+    if (this._isNetworkDisabled) {
+      return PROMPT_NAME_AFFINE_AI;
+    }
+    return this._isNetworkActive
+      ? PROMPT_NAME_NETWORK_SEARCH
+      : PROMPT_NAME_AFFINE_AI;
+  }
+
+  protected async updatePromptName(promptName: string) {
+    const sessionId = await this.getSessionId();
+    if (sessionId && AIProvider.session) {
+      await AIProvider.session.updateSession(sessionId, promptName);
+    }
+  }
+
+  override connectedCallback() {
+    super.connectedCallback();
+    this._disposables.add(
+      AIProvider.slots.requestSendWithChat.subscribe(
+        ({ input, context, host }) => {
+          if (this.host === host) {
+            context && this.updateContext(context);
+            const { updateComplete, send } = this;
+            updateComplete
+              .then(() => {
+                return send(input);
+              })
+              .catch(console.error);
+          }
+        }
+      )
+    );
+  }
+
+  protected override render() {
+    const { images, status } = this.chatContextValue;
+    const hasImages = images.length > 0;
+    const maxHeight = hasImages ? 272 + 2 : 200 + 2;
+    const uploadDisabled = this._isNetworkActive && !this._isNetworkDisabled;
+    return html` <div
+      class="chat-panel-input"
+      data-if-focused=${this.focused}
+      style=${styleMap({
+        maxHeight: `${maxHeight}px !important`,
+      })}
+      @pointerdown=${this._handlePointerDown}
+    >
+      ${hasImages
+        ? html`
+            <image-preview-grid
+              .images=${images}
+              .onImageRemove=${this._handleImageRemove}
+            ></image-preview-grid>
+          `
+        : nothing}
+      ${this.chatContextValue.quote
+        ? html`<div
+            class="chat-selection-quote"
+            data-testid="chat-selection-quote"
+          >
+            ${repeat(
+              getFirstTwoLines(this.chatContextValue.quote),
+              line => line,
+              line => html`<div>${line}</div>`
+            )}
+            <div
+              class="chat-quote-close"
+              @click=${() => {
+                this.updateContext({ quote: '', markdown: '' });
+              }}
+            >
+              ${CloseIcon()}
+            </div>
+          </div>`
+        : nothing}
+      <textarea
+        rows="1"
+        placeholder="What are your thoughts?"
+        @input=${this._handleInput}
+        @keydown=${this._handleKeyDown}
+        @focus=${() => {
+          this.focused = true;
+        }}
+        @blur=${() => {
+          this.focused = false;
+        }}
+        @paste=${this._handlePaste}
+        data-testid="chat-panel-input"
+      ></textarea>
+      <div class="chat-panel-input-actions">
+        <div
+          class="chat-history-clear"
+          aria-disabled=${this._isClearDisabled}
+          @click=${this._handleClear}
+          data-testid="chat-panel-clear"
+        >
+          ${BroomIcon()}
+        </div>
+        ${this.networkSearchConfig.visible.value
+          ? html`
+              <div
+                class="chat-network-search"
+                data-testid="chat-network-search"
+                aria-disabled=${this._isNetworkDisabled}
+                data-active=${this._isNetworkActive}
+                @click=${this._isNetworkDisabled
+                  ? undefined
+                  : this._toggleNetworkSearch}
+                @pointerdown=${stopPropagation}
+              >
+                ${PublishIcon()}
+              </div>
+            `
+          : nothing}
+        ${images.length < MaximumImageCount
+          ? html`<div
+              data-testid="chat-panel-input-image-upload"
+              class="image-upload"
+              aria-disabled=${uploadDisabled}
+              @click=${uploadDisabled ? undefined : this._uploadImageFiles}
+            >
+              ${ImageIcon()}
+            </div>`
+          : nothing}
+        ${status === 'transmitting'
+          ? html`<div @click=${this._handleAbort} data-testid="chat-panel-stop">
+              ${ChatAbortIcon}
+            </div>`
+          : html`<div
+              @click="${this._onTextareaSend}"
+              class="chat-panel-send"
+              aria-disabled=${this.isInputEmpty}
+              data-testid="chat-panel-send"
+            >
+              ${ChatSendIcon}
+            </div>`}
+      </div>
+    </div>`;
+  }
+
+  private readonly _handlePointerDown = (e: MouseEvent) => {
+    if (e.target !== this.textarea) {
+      // by default the div will be focused and will blur the textarea
+      e.preventDefault();
+      this.textarea.focus();
+    }
+  };
+
+  private readonly _handleInput = () => {
+    const { textarea } = this;
+    this.isInputEmpty = !textarea.value.trim();
+    textarea.style.height = 'auto';
+    textarea.style.height = textarea.scrollHeight + 'px';
+    let imagesHeight = this.imagePreviewGrid?.scrollHeight ?? 0;
+    if (imagesHeight) imagesHeight += 12;
+    if (this.scrollHeight >= 200 + imagesHeight) {
+      textarea.style.height = '148px';
+      textarea.style.overflowY = 'scroll';
+    }
+  };
+
+  private readonly _handleKeyDown = async (evt: KeyboardEvent) => {
+    if (evt.key === 'Enter' && !evt.shiftKey && !evt.isComposing) {
+      await this._onTextareaSend(evt);
+    }
+  };
+
+  private readonly _handlePaste = (event: ClipboardEvent) => {
+    const items = event.clipboardData?.items;
+    if (!items) return;
+
+    for (const index in items) {
+      const item = items[index];
+      if (item.kind === 'file' && item.type.indexOf('image') >= 0) {
+        const blob = item.getAsFile();
+        if (!blob) continue;
+        this._addImages([blob]);
+      }
+    }
+  };
+
+  private readonly _handleAbort = () => {
+    this.chatContextValue.abortController?.abort();
+    this.updateContext({ status: 'success' });
+    reportResponse('aborted:stop');
+  };
+
+  private readonly _handleClear = async () => {
+    if (this._isClearDisabled) {
+      return;
+    }
+    await this.cleanupHistories();
+  };
+
+  private readonly _toggleNetworkSearch = (e: MouseEvent) => {
+    e.preventDefault();
+    e.stopPropagation();
+
+    const enable = this.networkSearchConfig.enabled.value;
+    this.networkSearchConfig.setEnabled(!enable);
+  };
+
+  private _addImages(images: File[]) {
+    const oldImages = this.chatContextValue.images;
+    this.updateContext({
+      images: [...oldImages, ...images].slice(0, MaximumImageCount),
+    });
+  }
+
+  private readonly _handleImageRemove = (index: number) => {
+    const oldImages = this.chatContextValue.images;
+    const newImages = oldImages.filter((_, i) => i !== index);
+    this.updateContext({ images: newImages });
+  };
+
+  private readonly _uploadImageFiles = async (_e: MouseEvent) => {
+    const images = await openFileOrFiles({
+      acceptType: 'Images',
+      multiple: true,
+    });
+    if (!images) return;
+    this._addImages(images);
+  };
+
+  private readonly _onTextareaSend = async (e: MouseEvent | KeyboardEvent) => {
+    e.preventDefault();
+    e.stopPropagation();
+
+    const value = this.textarea.value.trim();
+    if (value.length === 0) return;
+
+    this.textarea.value = '';
+    this.isInputEmpty = true;
+    this.textarea.style.height = 'unset';
+
+    await this.send(value);
+  };
+
+  protected abstract send(text: string): Promise<void>;
+
+  protected async getMatchedContexts(userInput: string) {
+    const contextId = await this.getContextId();
+    if (!contextId) {
+      return { files: [], docs: [] };
+    }
+
+    const docContexts = new Map<
+      string,
+      { docId: string; docContent: string }
+    >();
+    const fileContexts = new Map<
+      string,
+      BlockSuitePresets.AIFileContextOption
+    >();
+
+    const { files: matchedFiles = [], docs: matchedDocs = [] } =
+      (await AIProvider.context?.matchContext(contextId, userInput)) ?? {};
+
+    matchedDocs.forEach(doc => {
+      docContexts.set(doc.docId, {
+        docId: doc.docId,
+        docContent: doc.content,
+      });
+    });
+
+    matchedFiles.forEach(file => {
+      const context = fileContexts.get(file.fileId);
+      if (context) {
+        context.fileContent += `\n${file.content}`;
+      } else {
+        const fileChip = this.chips.find(
+          chip => isFileChip(chip) && chip.fileId === file.fileId
+        ) as FileChip | undefined;
+        if (fileChip && fileChip.blobId) {
+          fileContexts.set(file.fileId, {
+            blobId: fileChip.blobId,
+            fileName: fileChip.file.name,
+            fileType: fileChip.file.type,
+            fileContent: file.content,
+          });
+        }
+      }
+    });
+
+    this.chips.forEach(chip => {
+      if (isDocChip(chip) && !!chip.markdown?.value) {
+        docContexts.set(chip.docId, {
+          docId: chip.docId,
+          docContent: chip.markdown.value,
+        });
+      }
+    });
+
+    const docs: BlockSuitePresets.AIDocContextOption[] = Array.from(
+      docContexts.values()
+    ).map(doc => {
+      const docMeta = this.docDisplayConfig.getDocMeta(doc.docId);
+      const docTitle = this.docDisplayConfig.getTitle(doc.docId);
+      const tags = docMeta?.tags
+        ? docMeta.tags
+            .map(tagId => this.docDisplayConfig.getTagTitle(tagId))
+            .join(',')
+        : '';
+      return {
+        docId: doc.docId,
+        docContent: doc.docContent,
+        docTitle,
+        tags,
+        createDate: docMeta?.createDate
+          ? new Date(docMeta.createDate).toISOString()
+          : '',
+        updatedDate: docMeta?.updatedDate
+          ? new Date(docMeta.updatedDate).toISOString()
+          : '',
+      };
+    });
+
+    return {
+      docs,
+      files: Array.from(fileContexts.values()),
+    };
+  }
+}
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/const.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/const.ts
@@ -0,0 +1,2 @@
+export const PROMPT_NAME_AFFINE_AI = 'Chat With AFFiNE AI';
+export const PROMPT_NAME_NETWORK_SEARCH = 'Search With AFFiNE AI';
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/index.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/index.ts
@@ -0,0 +1,3 @@
+export * from './ai-chat-input';
+export * from './const';
+export * from './type';
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/type.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/type.ts
@@ -0,0 +1,21 @@
+import type { Signal } from '@preact/signals-core';
+
+import type { AIError } from '../../provider';
+import type { ChatStatus, HistoryMessage } from '../ai-chat-messages';
+
+export interface AINetworkSearchConfig {
+  visible: Signal<boolean | undefined>;
+  enabled: Signal<boolean | undefined>;
+  setEnabled: (state: boolean) => void;
+}
+
+// TODO: remove this type
+export type AIChatInputContext = {
+  messages: HistoryMessage[];
+  status: ChatStatus;
+  error: AIError | null;
+  quote?: string;
+  markdown?: string;
+  images: File[];
+  abortController: AbortController | null;
+};
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-messages/index.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-messages/index.ts
@@ -0,0 +1 @@
+export * from './type';
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-messages/type.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-messages/type.ts
@@ -1,6 +1,5 @@
 import { z } from 'zod';

-// Define the Zod schema
 const ChatMessageSchema = z.object({
  id: z.string(),
  content: z.string(),
@@ -14,12 +13,36 @@ const ChatMessageSchema = z.object({

 export const ChatMessagesSchema = z.array(ChatMessageSchema);

-// Derive the TypeScript type from the Zod schema
 export type ChatMessage = z.infer<typeof ChatMessageSchema>;

+export type ChatAction = {
+  action: string;
+  messages: ChatMessage[];
+  sessionId: string;
+  createdAt: string;
+};
+
+export type HistoryMessage = ChatMessage | ChatAction;
+
 export type MessageRole = 'user' | 'assistant';
+
 export type MessageUserInfo = {
  userId?: string;
  userName?: string;
  avatarUrl?: string;
 };
+
+export function isChatAction(item: HistoryMessage): item is ChatAction {
+  return 'action' in item;
+}
+
+export function isChatMessage(item: HistoryMessage): item is ChatMessage {
+  return 'role' in item;
+}
+
+export type ChatStatus =
+  | 'loading'
+  | 'success'
+  | 'error'
+  | 'idle'
+  | 'transmitting';
--- a/packages/frontend/core/src/blocksuite/ai/peek-view/chat-block-input.ts
+++ b/packages/frontend/core/src/blocksuite/ai/peek-view/chat-block-input.ts
@@ -1,229 +1,13 @@
-import { SignalWatcher } from '@blocksuite/affine/global/lit';
-import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
-import { openFileOrFiles } from '@blocksuite/affine/shared/utils';
-import type { EditorHost } from '@blocksuite/affine/std';
-import { BroomIcon, ImageIcon, PublishIcon } from '@blocksuite/icons/lit';
-import { css, html, LitElement, nothing } from 'lit';
-import { property, query, state } from 'lit/decorators.js';
-import { classMap } from 'lit/directives/class-map.js';
+import { property } from 'lit/decorators.js';

-import { ChatAbortIcon, ChatSendIcon } from '../_common/icons';
-import type { ChatMessage } from '../blocks';
-import type { AINetworkSearchConfig } from '../chat-panel/chat-config';
-import {
-  PROMPT_NAME_AFFINE_AI,
-  PROMPT_NAME_NETWORK_SEARCH,
-} from '../chat-panel/const';
+import { AIChatInput } from '../components/ai-chat-input';
+import type { ChatMessage } from '../components/ai-chat-messages';
 import { type AIError, AIProvider } from '../provider';
-import { reportResponse } from '../utils/action-reporter';
 import { readBlobAsURL } from '../utils/image';
-import { stopPropagation } from '../utils/selection-utils';
-import type { ChatContext } from './types';
-
-const MaximumImageCount = 8;
-
-export class ChatBlockInput extends SignalWatcher(LitElement) {
-  static override styles = css`
-    :host {
-      width: 100%;
-    }
-    .ai-chat-input {
-      display: flex;
-      width: 100%;
-      min-height: 100px;
-      max-height: 206px;
-      padding: 8px 12px;
-      box-sizing: border-box;
-      border: 1px solid var(--affine-border-color);
-      border-radius: 4px;
-      flex-direction: column;
-      justify-content: space-between;
-      gap: 12px;
-      position: relative;
-      background-color: var(--affine-white-10);
-    }
-    .ai-chat-input {
-      textarea {
-        width: 100%;
-        padding: 0;
-        margin: 0;
-        border: none;
-        line-height: 22px;
-        font-size: var(--affine-font-sm);
-        font-weight: 400;
-        font-family: var(--affine-font-family);
-        color: var(--affine-text-primary-color);
-        box-sizing: border-box;
-        resize: none;
-        overflow-y: hidden;
-        background-color: transparent;
-        user-select: none;
-      }
-      textarea::placeholder {
-        font-size: 14px;
-        font-weight: 400;
-        font-family: var(--affine-font-family);
-        color: var(--affine-placeholder-color);
-      }
-      textarea:focus {
-        outline: none;
-      }
-    }
-
-    .chat-panel-send svg rect {
-      fill: var(--affine-primary-color);
-    }
-    .chat-panel-send[aria-disabled='true'] {
-      cursor: not-allowed;
-    }
-    .chat-panel-send[aria-disabled='true'] svg rect {
-      fill: var(--affine-text-disable-color);
-    }
-
-    .chat-panel-input-actions {
-      display: flex;
-      gap: 8px;
-      align-items: center;
-
-      div {
-        width: 24px;
-        height: 24px;
-        cursor: pointer;
-      }
-      div:nth-child(2) {
-        margin-left: auto;
-      }
-
-      .image-upload,
-      .chat-history-clear,
-      .chat-network-search {
-        display: flex;
-        justify-content: center;
-        align-items: center;
-        svg {
-          width: 20px;
-          height: 20px;
-          color: ${unsafeCSSVarV2('icon/primary')};
-        }
-      }
-      .chat-network-search[data-active='true'] svg {
-        color: ${unsafeCSSVarV2('icon/activated')};
-      }
-
-      .chat-network-search[aria-disabled='true'] {
-        cursor: not-allowed;
-      }
-      .chat-network-search[aria-disabled='true'] svg {
-        color: var(--affine-text-disable-color) !important;
-      }
-    }
-
-    .chat-history-clear.disabled {
-      cursor: not-allowed;
-      opacity: 0.5;
-    }
-  `;
-
-  override render() {
-    const { images, status, messages } = this.chatContext;
-    const hasImages = images.length > 0;
-    const maxHeight = hasImages ? 272 + 2 : 200 + 2;
-    const disableCleanUp =
-      status === 'loading' || status === 'transmitting' || !messages.length;
-    const cleanButtonClasses = classMap({
-      'chat-history-clear': true,
-      disabled: disableCleanUp,
-    });
-
-    return html`<style>
-        .chat-panel-input {
-          border-color: ${this._focused
-            ? 'var(--affine-primary-color)'
-            : 'var(--affine-border-color)'};
-          box-shadow: ${this._focused ? 'var(--affine-active-shadow)' : 'none'};
-          max-height: ${maxHeight}px;
-          user-select: none;
-        }
-      </style>
-      <div class="ai-chat-input">
-        ${hasImages
-          ? html`<image-preview-grid
-              .images=${images}
-              .onImageRemove=${this._handleImageRemove}
-            ></image-preview-grid>`
-          : nothing}
-        <textarea
-          rows="1"
-          placeholder="What are your thoughts?"
-          @keydown=${this._handleKeyDown}
-          @input=${this._handleInput}
-          @focus=${() => {
-            this._focused = true;
-          }}
-          @blur=${() => {
-            this._focused = false;
-          }}
-          @paste=${this._handlePaste}
-          data-testid="chat-block-input"
-        ></textarea>
-        <div class="chat-panel-input-actions">
-          <div class=${cleanButtonClasses} @click=${this._handleCleanup}>
-            ${BroomIcon()}
-          </div>
-          ${this.networkSearchConfig.visible.value
-            ? html`
-                <div
-                  class="chat-network-search"
-                  data-testid="chat-network-search"
-                  aria-disabled=${this._isNetworkDisabled}
-                  data-active=${this._isNetworkActive}
-                  @click=${this._isNetworkDisabled
-                    ? undefined
-                    : this._toggleNetworkSearch}
-                  @pointerdown=${stopPropagation}
-                >
-                  ${PublishIcon()}
-                </div>
-              `
-            : nothing}
-          ${images.length < MaximumImageCount
-            ? html`<div
-                data-testid="chat-block-input-image-upload"
-                class="image-upload"
-                @click=${this._handleImageUpload}
-              >
-                ${ImageIcon()}
-              </div>`
-            : nothing}
-          ${status === 'transmitting'
-            ? html`<div
-                @click=${this._handleAbort}
-                data-testid="chat-panel-peek-view-stop"
-              >
-                ${ChatAbortIcon}
-              </div>`
-            : html`<div
-                @click=${this._onTextareaSend}
-                class="chat-panel-send"
-                aria-disabled=${this._isInputEmpty}
-              >
-                ${ChatSendIcon}
-              </div>`}
-        </div>
-      </div>`;
-  }

+export class ChatBlockInput extends AIChatInput {
  @property({ attribute: false })
-  accessor parentSessionId!: string;
-
-  @property({ attribute: false })
-  accessor latestMessageId!: string;
-
-  @property({ attribute: false })
-  accessor host!: EditorHost;
-
-  @property({ attribute: false })
-  accessor networkSearchConfig!: AINetworkSearchConfig;
+  accessor getBlockId!: () => string | null | undefined;

  @property({ attribute: false })
  accessor updateChatBlock!: () => Promise<void>;
@@ -231,147 +15,10 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {
  @property({ attribute: false })
  accessor createChatBlock!: () => Promise<void>;

-  @property({ attribute: false })
-  accessor cleanupHistories!: () => Promise<void>;
-
-  @property({ attribute: false })
-  accessor updateContext!: (context: Partial<ChatContext>) => void;
-
-  @property({ attribute: false })
-  accessor chatContext!: ChatContext;
-
-  @query('textarea')
-  accessor textarea!: HTMLTextAreaElement;
-
-  @state()
-  accessor _isInputEmpty = true;
-
-  @state()
-  accessor _focused = false;
-
-  private get _isNetworkActive() {
-    return (
-      !!this.networkSearchConfig.visible.value &&
-      !!this.networkSearchConfig.enabled.value
-    );
-  }
-
-  private get _isNetworkDisabled() {
-    return !!this.chatContext.images.length;
-  }
-
-  private _getPromptName() {
-    if (this._isNetworkDisabled) {
-      return PROMPT_NAME_AFFINE_AI;
-    }
-    return this._isNetworkActive
-      ? PROMPT_NAME_NETWORK_SEARCH
-      : PROMPT_NAME_AFFINE_AI;
-  }
-
-  private async _updatePromptName(promptName: string) {
-    const { currentSessionId } = this.chatContext;
-    if (currentSessionId && AIProvider.session) {
-      await AIProvider.session.updateSession(currentSessionId, promptName);
-    }
-  }
-
-  private readonly _addImages = (images: File[]) => {
-    const oldImages = this.chatContext.images;
-    this.updateContext({
-      images: [...oldImages, ...images].slice(0, MaximumImageCount),
-    });
-  };
-
-  private readonly _toggleNetworkSearch = (e: MouseEvent) => {
-    e.preventDefault();
-    e.stopPropagation();
-
-    const enable = this.networkSearchConfig.enabled.value;
-    this.networkSearchConfig.setEnabled(!enable);
-  };
-
-  private readonly _handleKeyDown = async (evt: KeyboardEvent) => {
-    if (evt.key === 'Enter' && !evt.shiftKey && !evt.isComposing) {
-      evt.preventDefault();
-      await this._onTextareaSend(evt);
-    }
-  };
-
-  private readonly _handleInput = () => {
-    const { textarea } = this;
-    this._isInputEmpty = !textarea.value.trim();
-    textarea.style.height = 'auto';
-    textarea.style.height = textarea.scrollHeight + 'px';
-    if (this.scrollHeight >= 202) {
-      textarea.style.height = '168px';
-      textarea.style.overflowY = 'scroll';
-    }
-  };
-
-  private readonly _handlePaste = (event: ClipboardEvent) => {
-    const items = event.clipboardData?.items;
-    if (!items) return;
-    for (const index in items) {
-      const item = items[index];
-      if (item.kind === 'file' && item.type.indexOf('image') >= 0) {
-        const blob = item.getAsFile();
-        if (!blob) continue;
-        this._addImages([blob]);
-      }
-    }
-  };
-
-  private readonly _handleCleanup = async () => {
-    if (
-      this.chatContext.status === 'loading' ||
-      this.chatContext.status === 'transmitting' ||
-      !this.chatContext.messages.length
-    ) {
-      return;
-    }
-    await this.cleanupHistories();
-  };
-
-  private readonly _handleImageUpload = async () => {
-    const images = await openFileOrFiles({
-      acceptType: 'Images',
-      multiple: true,
-    });
-    if (!images) return;
-    this._addImages(images);
-  };
-
-  private readonly _handleAbort = () => {
-    this.chatContext.abortController?.abort();
-    this.updateContext({ status: 'success' });
-    reportResponse('aborted:stop');
-  };
-
-  private readonly _handleImageRemove = (index: number) => {
-    const oldImages = this.chatContext.images;
-    const newImages = oldImages.filter((_, i) => i !== index);
-    this.updateContext({ images: newImages });
-  };
-
-  private readonly _onTextareaSend = async (e: MouseEvent | KeyboardEvent) => {
-    e.preventDefault();
-    e.stopPropagation();
-
-    const value = this.textarea.value.trim();
-    if (value.length === 0) return;
-
-    this.textarea.value = '';
-    this._isInputEmpty = true;
-    this.textarea.style.height = 'unset';
-
-    await this._send(value);
-  };
-
-  private readonly _send = async (text: string) => {
-    const { images, status, currentChatBlockId, currentSessionId } =
-      this.chatContext;
-    const chatBlockExists = !!currentChatBlockId;
+  send = async (text: string) => {
+    const { images, status } = this.chatContextValue;
+    const sessionId = await this.getSessionId();
+    if (!sessionId) return;
    let content = '';

    if (status === 'loading' || status === 'transmitting') return;
@@ -379,7 +26,7 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {

    try {
      const { doc } = this.host;
-      const promptName = this._getPromptName();
+      const promptName = this.getPromptName();

      this.updateContext({
        images: [],
@@ -394,7 +41,7 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {
      const userInfo = await AIProvider.userInfo;
      this.updateContext({
        messages: [
-          ...this.chatContext.messages,
+          ...this.chatContextValue.messages,
          {
            id: '',
            content: text,
@@ -416,28 +63,12 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {

      // must update prompt name after local chat message is updated
      // otherwise, the unauthorized error can not be rendered properly
-      await this._updatePromptName(promptName);
-
-      // If has not forked a chat session, fork a new one
-      let chatSessionId = currentSessionId;
-      if (!chatSessionId) {
-        const forkSessionId = await AIProvider.forkChat?.({
-          workspaceId: doc.workspace.id,
-          docId: doc.id,
-          sessionId: this.parentSessionId,
-          latestMessageId: this.latestMessageId,
-        });
-        if (!forkSessionId) return;
-        this.updateContext({
-          currentSessionId: forkSessionId,
-        });
-        chatSessionId = forkSessionId;
-      }
+      await this.updatePromptName(promptName);

      const abortController = new AbortController();
      const stream = AIProvider.actions.chat?.({
        input: text,
-        sessionId: chatSessionId,
+        sessionId,
        docId: doc.id,
        attachments: images,
        workspaceId: doc.workspace.id,
@@ -454,7 +85,7 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {
        });

        for await (const text of stream) {
-          const messages = [...this.chatContext.messages];
+          const messages = [...this.chatContextValue.messages];
          const last = messages[messages.length - 1] as ChatMessage;
          last.content += text;
          this.updateContext({ messages, status: 'transmitting' });
@@ -468,6 +99,7 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {
      this.updateContext({ status: 'error', error: error as AIError });
    } finally {
      if (content) {
+        const chatBlockExists = !!this.getBlockId();
        if (!chatBlockExists) {
          await this.createChatBlock();
        }
--- a/packages/frontend/core/src/blocksuite/ai/peek-view/chat-block-peek-view.ts
+++ b/packages/frontend/core/src/blocksuite/ai/peek-view/chat-block-peek-view.ts
@@ -22,12 +22,10 @@ import {
  constructUserInfoWithMessages,
  queryHistoryMessages,
 } from '../_common/chat-actions-handle';
-import {
-  type AIChatBlockModel,
-  type ChatMessage,
-  ChatMessagesSchema,
-} from '../blocks';
-import type { AINetworkSearchConfig } from '../chat-panel/chat-config';
+import { type AIChatBlockModel } from '../blocks';
+import type { AINetworkSearchConfig } from '../components/ai-chat-input';
+import type { ChatMessage } from '../components/ai-chat-messages';
+import { ChatMessagesSchema } from '../components/ai-chat-messages';
 import type { TextRendererOptions } from '../components/text-renderer';
 import { AIChatErrorRenderer } from '../messages/error';
 import { type AIError, AIProvider } from '../provider';
@@ -64,6 +62,12 @@ export class AIChatBlockPeekView extends LitElement {

  private _textRendererOptions: TextRendererOptions = {};

+  private _chatSessionId: string | null | undefined = null;
+
+  private _chatContextId: string | null | undefined = null;
+
+  private _chatBlockId: string | null | undefined = null;
+
  private readonly _deserializeHistoryChatMessages = (
    historyMessagesString: string
  ) => {
@@ -127,13 +131,35 @@ export class AIChatBlockPeekView extends LitElement {
      images: [],
      abortController: null,
      messages: [],
-      currentSessionId: null,
-      currentChatBlockId: null,
    });
+    this._chatSessionId = null;
+    this._chatContextId = null;
+    this._chatBlockId = null;
  };

  private readonly _getSessionId = async () => {
-    return this.chatContext.currentSessionId ?? undefined;
+    // If has not forked a chat session, fork a new one
+    if (!this._chatSessionId) {
+      const latestMessage = this._historyMessages.at(-1);
+      if (!latestMessage) return;
+
+      const forkSessionId = await AIProvider.forkChat?.({
+        workspaceId: this.host.doc.workspace.id,
+        docId: this.host.doc.id,
+        sessionId: this.parentSessionId,
+        latestMessageId: latestMessage.id,
+      });
+      this._chatSessionId = forkSessionId;
+    }
+    return this._chatSessionId;
+  };
+
+  private readonly _getContextId = async () => {
+    return this._chatContextId;
+  };
+
+  private readonly _getBlockId = () => {
+    return this._chatBlockId;
  };

  /**
@@ -147,15 +173,12 @@ export class AIChatBlockPeekView extends LitElement {
    }

    // If there is already a chat block, do not create a new one
-    if (this.chatContext.currentChatBlockId) {
+    if (this._chatBlockId) {
      return;
    }

    // If there is no session id or chat messages, do not create a new chat block
-    if (
-      !this.chatContext.currentSessionId ||
-      !this.chatContext.messages.length
-    ) {
+    if (!this._chatSessionId || !this.chatContext.messages.length) {
      return;
    }

@@ -173,7 +196,7 @@ export class AIChatBlockPeekView extends LitElement {
    const messages = await this._constructBranchChatBlockMessages(
      parentRootWorkspaceId,
      parentRootDocId,
-      this.chatContext.currentSessionId
+      this._chatSessionId
    );
    if (!messages.length) {
      return;
@@ -187,7 +210,7 @@ export class AIChatBlockPeekView extends LitElement {
      {
        xywh: bound.serialize(),
        messages: JSON.stringify(messages),
-        sessionId: this.chatContext.currentSessionId,
+        sessionId: this._chatSessionId,
        rootWorkspaceId: parentRootWorkspaceId,
        rootDocId: parentRootDocId,
      },
@@ -198,7 +221,7 @@ export class AIChatBlockPeekView extends LitElement {
      return;
    }

-    this.updateContext({ currentChatBlockId: aiChatBlockId });
+    this._chatBlockId = aiChatBlockId;

    // Connect the parent chat block to the AI chat block
    crud.addElement(CanvasElementType.CONNECTOR, {
@@ -223,15 +246,12 @@ export class AIChatBlockPeekView extends LitElement {
   * Update the current chat messages with the new message
   */
  updateChatBlockMessages = async () => {
-    if (
-      !this.chatContext.currentChatBlockId ||
-      !this.chatContext.currentSessionId
-    ) {
+    if (!this._chatBlockId || !this._chatSessionId) {
      return;
    }

    const { doc } = this.host;
-    const chatBlock = doc.getBlock(this.chatContext.currentChatBlockId);
+    const chatBlock = doc.getBlock(this._chatBlockId);
    if (!chatBlock) return;

    // Get fork session messages
@@ -239,7 +259,7 @@ export class AIChatBlockPeekView extends LitElement {
    const messages = await this._constructBranchChatBlockMessages(
      parentRootWorkspaceId,
      parentRootDocId,
-      this.chatContext.currentSessionId
+      this._chatSessionId
    );
    if (!messages.length) {
      return;
@@ -260,8 +280,8 @@ export class AIChatBlockPeekView extends LitElement {
    const notificationService = this.host.std.getOptional(NotificationProvider);
    if (!notificationService) return;

-    const { currentChatBlockId, currentSessionId } = this.chatContext;
-    if (!currentChatBlockId && !currentSessionId) {
+    const { _chatBlockId, _chatSessionId } = this;
+    if (!_chatBlockId && !_chatSessionId) {
      return;
    }

@@ -275,21 +295,21 @@ export class AIChatBlockPeekView extends LitElement {
      })
    ) {
      const { doc } = this.host;
-      if (currentSessionId) {
+      if (_chatSessionId) {
        await AIProvider.histories?.cleanup(doc.workspace.id, doc.id, [
-          currentSessionId,
+          _chatSessionId,
        ]);
      }

-      if (currentChatBlockId) {
+      if (_chatBlockId) {
        const surface = getSurfaceBlock(doc);
        const crud = this.host.std.get(EdgelessCRUDIdentifier);
-        const chatBlock = doc.getBlock(currentChatBlockId)?.model;
+        const chatBlock = doc.getBlock(_chatBlockId)?.model;
        if (chatBlock) {
          const connectors = surface?.getConnectors(chatBlock.id);
          doc.transact(() => {
            // Delete the AI chat block
-            crud.removeElement(currentChatBlockId);
+            crud.removeElement(_chatBlockId);
            // Delete the connectors
            connectors?.forEach(connector => {
              crud.removeElement(connector.id);
@@ -308,8 +328,8 @@ export class AIChatBlockPeekView extends LitElement {
   */
  retry = async () => {
    const { doc } = this.host;
-    const { currentChatBlockId, currentSessionId } = this.chatContext;
-    if (!currentChatBlockId || !currentSessionId) {
+    const { _chatBlockId, _chatSessionId } = this;
+    if (!_chatBlockId || !_chatSessionId) {
      return;
    }

@@ -326,7 +346,7 @@ export class AIChatBlockPeekView extends LitElement {
      this.updateContext({ messages, status: 'loading', error: null });

      const stream = AIProvider.actions.chat?.({
-        sessionId: currentSessionId,
+        sessionId: _chatSessionId,
        retry: true,
        docId: doc.id,
        workspaceId: doc.workspace.id,
@@ -478,9 +498,7 @@ export class AIChatBlockPeekView extends LitElement {

    const latestHistoryMessage = _historyMessages[_historyMessages.length - 1];
    const latestMessageCreatedAt = latestHistoryMessage.createdAt;
-    const latestHistoryMessageId = latestHistoryMessage.id;
    const {
-      parentSessionId,
      updateChatBlockMessages,
      createAIChatBlock,
      cleanCurrentChatHistories,
@@ -506,12 +524,13 @@ export class AIChatBlockPeekView extends LitElement {
      </div>
      <chat-block-input
        .host=${host}
-        .parentSessionId=${parentSessionId}
-        .latestMessageId=${latestHistoryMessageId}
+        .getSessionId=${this._getSessionId}
+        .getContextId=${this._getContextId}
+        .getBlockId=${this._getBlockId}
        .updateChatBlock=${updateChatBlockMessages}
        .createChatBlock=${createAIChatBlock}
        .cleanupHistories=${cleanCurrentChatHistories}
-        .chatContext=${chatContext}
+        .chatContextValue=${chatContext}
        .updateContext=${updateContext}
        .networkSearchConfig=${networkSearchConfig}
      ></chat-block-input>
@@ -547,8 +566,6 @@ export class AIChatBlockPeekView extends LitElement {
    images: [],
    abortController: null,
    messages: [],
-    currentSessionId: null,
-    currentChatBlockId: null,
  };
 }

--- a/packages/frontend/core/src/blocksuite/ai/peek-view/types.ts
+++ b/packages/frontend/core/src/blocksuite/ai/peek-view/types.ts
@@ -1,19 +1,10 @@
-import type { ChatMessage } from '../blocks';
+import type { ChatMessage, ChatStatus } from '../components/ai-chat-messages';
 import type { AIError } from '../provider';

-export type ChatStatus =
-  | 'success'
-  | 'error'
-  | 'idle'
-  | 'transmitting'
-  | 'loading';
-
 export type ChatContext = {
  messages: ChatMessage[];
  status: ChatStatus;
  error: AIError | null;
  images: File[];
  abortController: AbortController | null;
-  currentSessionId: string | null;
-  currentChatBlockId: string | null;
 };