diff --git 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
deleted file mode 100644
index d970b81a6a..0000000000
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-input.ts
+++ /dev/null
@@ -1,108 +0,0 @@
-import { AIChatInput } from '../components/ai-chat-input';
-import type { ChatMessage } from '../components/ai-chat-messages';
-import { type AIError, AIProvider } from '../provider';
-import { readBlobAsURL } from '../utils/image';
-
-export class ChatPanelInput extends AIChatInput {
-  send = async (text: string) => {
-    const { status, markdown, images } = this.chatContextValue;
-    if (status === 'loading' || status === 'transmitting') return;
-    if (!text) return;
-
-    try {
-      const { doc } = this.host;
-      const promptName = this.getPromptName();
-
-      this.updateContext({
-        images: [],
-        status: 'loading',
-        error: null,
-        quote: '',
-        markdown: '',
-      });
-
-      const attachments = await Promise.all(
-        images?.map(image => readBlobAsURL(image))
-      );
-
-      const userInput = (markdown ? `${markdown}\n` : '') + text;
-      this.updateContext({
-        messages: [
-          ...this.chatContextValue.messages,
-          {
-            id: '',
-            role: 'user',
-            content: userInput,
-            createdAt: new Date().toISOString(),
-            attachments,
-          },
-          {
-            id: '',
-            role: 'assistant',
-            content: '',
-            createdAt: new Date().toISOString(),
-          },
-        ],
-      });
-
-      // must update prompt name after local chat message is updated
-      // otherwise, the unauthorized error can not be rendered properly
-      await this.updatePromptName(promptName);
-
-      const abortController = new AbortController();
-      const sessionId = await this.getSessionId();
-      if (!sessionId) return;
-
-      const contexts = await this.getMatchedContexts(userInput);
-      const stream = AIProvider.actions.chat?.({
-        sessionId,
-        input: userInput,
-        contexts,
-        docId: doc.id,
-        attachments: images,
-        workspaceId: doc.workspace.id,
-        host: this.host,
-        stream: true,
-        signal: abortController.signal,
-        where: 'chat-panel',
-        control: 'chat-send',
-        isRootSession: true,
-      });
-
-      if (stream) {
-        this.updateContext({ abortController });
-
-        for await (const text of stream) {
-          const messages = [...this.chatContextValue.messages];
-          const last = messages[messages.length - 1] as ChatMessage;
-          last.content += text;
-          this.updateContext({ messages, status: 'transmitting' });
-        }
-
-        this.updateContext({ status: 'success' });
-
-        const { messages } = this.chatContextValue;
-        const last = messages[messages.length - 1] as ChatMessage;
-        if (!last.id) {
-          const historyIds = await AIProvider.histories?.ids(
-            doc.workspace.id,
-            doc.id,
-            { sessionId }
-          );
-          if (!historyIds || !historyIds[0]) return;
-          last.id = historyIds[0].messages.at(-1)?.id ?? '';
-        }
-      }
-    } catch (error) {
-      this.updateContext({ status: 'error', error: error as AIError });
-    } finally {
-      this.updateContext({ abortController: null });
-    }
-  };
-}
-
-declare global {
-  interface HTMLElementTagNameMap {
-    'chat-panel-input': ChatPanelInput;
-  }
-}
diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
index 616fd417d0..6926832e3c 100644
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
@@ -1,4 +1,3 @@
-import './chat-panel-input';
 import './chat-panel-messages';
 
 import type {
@@ -640,7 +639,7 @@ export class ChatPanel extends SignalWatcher(
         .docDisplayConfig=${this.docDisplayConfig}
         .searchMenuConfig=${this.searchMenuConfig}
       ></chat-panel-chips>
-      <chat-panel-input
+      <ai-chat-input
         .chips=${this.chips}
         .chatContextValue=${this.chatContextValue}
         .getSessionId=${this._getSessionId}
@@ -650,7 +649,7 @@ export class ChatPanel extends SignalWatcher(
         .updateContext=${this.updateContext}
         .host=${this.host}
         .cleanupHistories=${this._cleanupHistories}
-      ></chat-panel-input>
+      ></ai-chat-input>
       <div class="chat-panel-footer">
         ${InformationIcon()}
         <div>AI outputs can be misleading or wrong</div>
diff --git 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
index c9a0a0d667..d09a90012d 100644
--- 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
@@ -15,14 +15,16 @@ 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 { type AIError, AIProvider } from '../../provider';
 import { reportResponse } from '../../utils/action-reporter';
+import { readBlobAsURL } from '../../utils/image';
 import type {
   ChatChip,
   DocDisplayConfig,
   FileChip,
 } from '../ai-chat-chips/type';
 import { isDocChip, isFileChip } from '../ai-chat-chips/utils';
+import type { ChatMessage } from '../ai-chat-messages';
 import { PROMPT_NAME_AFFINE_AI, PROMPT_NAME_NETWORK_SEARCH } from './const';
 import type { AIChatInputContext, AINetworkSearchConfig } from './type';
 
@@ -33,9 +35,7 @@ function getFirstTwoLines(text: string) {
   return lines.slice(0, 2);
 }
 
-export abstract class AIChatInput extends SignalWatcher(
-  WithDisposable(LitElement)
-) {
+export class AIChatInput extends SignalWatcher(WithDisposable(LitElement)) {
   static override styles = css`
     :host {
       width: 100%;
@@ -220,10 +220,10 @@ export abstract class AIChatInput extends SignalWatcher(
   accessor chips: ChatChip[] = [];
 
   @property({ attribute: false })
-  accessor getSessionId!: () => Promise<string | null | undefined>;
+  accessor getSessionId!: () => Promise<string | undefined>;
 
   @property({ attribute: false })
-  accessor getContextId!: () => Promise<string | null | undefined>;
+  accessor getContextId!: () => Promise<string | undefined>;
 
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<AIChatInputContext>) => void;
@@ -237,6 +237,18 @@ export abstract class AIChatInput extends SignalWatcher(
   @property({ attribute: false })
   accessor docDisplayConfig!: DocDisplayConfig;
 
+  @property({ attribute: false })
+  accessor isRootSession: boolean = true;
+
+  @property({ attribute: false })
+  accessor onChatSuccess = () => null;
+
+  @property({ attribute: false })
+  accessor trackOptions: BlockSuitePresets.TrackerOptions = {
+    control: 'chat-send',
+    where: 'chat-panel',
+  };
+
   @property({ attribute: 'data-testid', reflect: true })
   accessor testId = 'chat-panel-input-container';
 
@@ -262,7 +274,7 @@ export abstract class AIChatInput extends SignalWatcher(
     );
   }
 
-  protected getPromptName() {
+  private _getPromptName() {
     if (this._isNetworkDisabled) {
       return PROMPT_NAME_AFFINE_AI;
     }
@@ -271,7 +283,7 @@ export abstract class AIChatInput extends SignalWatcher(
       : PROMPT_NAME_AFFINE_AI;
   }
 
-  protected async updatePromptName(promptName: string) {
+  private async _updatePromptName(promptName: string) {
     const sessionId = await this.getSessionId();
     if (sessionId && AIProvider.session) {
       await AIProvider.session.updateSession(sessionId, promptName);
@@ -501,9 +513,117 @@ export abstract class AIChatInput extends SignalWatcher(
     await this.send(value);
   };
 
-  protected abstract send(text: string): Promise<void>;
+  send = async (text: string) => {
+    const { status, markdown, images } = this.chatContextValue;
+    if (status === 'loading' || status === 'transmitting') return;
+    if (!text) return;
+    if (!AIProvider.actions.chat) return;
 
-  protected async getMatchedContexts(userInput: string) {
+    try {
+      const promptName = this._getPromptName();
+      const abortController = new AbortController();
+      this.updateContext({
+        images: [],
+        status: 'loading',
+        error: null,
+        quote: '',
+        markdown: '',
+        abortController,
+      });
+
+      const attachments = await Promise.all(
+        images?.map(image => readBlobAsURL(image))
+      );
+      const userInput = (markdown ? `${markdown}\n` : '') + text;
+
+      // optimistic update messages
+      await this._preUpdateMessages(userInput, attachments);
+      // must update prompt name after local chat message is updated
+      // otherwise, the unauthorized error can not be rendered properly
+      await this._updatePromptName(promptName);
+
+      const sessionId = await this.getSessionId();
+      const contexts = await this._getMatchedContexts(userInput);
+      if (abortController.signal.aborted) {
+        return;
+      }
+      const stream = AIProvider.actions.chat({
+        sessionId,
+        input: userInput,
+        contexts,
+        docId: this.host.doc.id,
+        attachments: images,
+        workspaceId: this.host.doc.workspace.id,
+        host: this.host,
+        stream: true,
+        signal: abortController.signal,
+        isRootSession: this.isRootSession,
+        where: this.trackOptions.where,
+        control: this.trackOptions.control,
+      });
+
+      for await (const text of stream) {
+        const messages = [...this.chatContextValue.messages];
+        const last = messages[messages.length - 1] as ChatMessage;
+        last.content += text;
+        this.updateContext({ messages, status: 'transmitting' });
+      }
+
+      this.updateContext({ status: 'success' });
+      this.onChatSuccess();
+      // update message id from server
+      await this._postUpdateMessages();
+    } catch (error) {
+      this.updateContext({ status: 'error', error: error as AIError });
+    } finally {
+      this.updateContext({ abortController: null });
+    }
+  };
+
+  private readonly _preUpdateMessages = async (
+    userInput: string,
+    attachments: string[]
+  ) => {
+    const userInfo = await AIProvider.userInfo;
+    this.updateContext({
+      messages: [
+        ...this.chatContextValue.messages,
+        {
+          id: '',
+          role: 'user',
+          content: userInput,
+          createdAt: new Date().toISOString(),
+          attachments,
+          userId: userInfo?.id,
+          userName: userInfo?.name,
+          avatarUrl: userInfo?.avatarUrl ?? undefined,
+        },
+        {
+          id: '',
+          role: 'assistant',
+          content: '',
+          createdAt: new Date().toISOString(),
+        },
+      ],
+    });
+  };
+
+  private readonly _postUpdateMessages = async () => {
+    const { messages } = this.chatContextValue;
+    const last = messages[messages.length - 1] as ChatMessage;
+    if (!last.id) {
+      const sessionId = await this.getSessionId();
+      const historyIds = await AIProvider.histories?.ids(
+        this.host.doc.workspace.id,
+        this.host.doc.id,
+        { sessionId }
+      );
+      if (!historyIds || !historyIds[0]) return;
+      last.id = historyIds[0].messages.at(-1)?.id ?? '';
+    }
+  };
+
+  private async _getMatchedContexts(userInput: string) {
     const contextId = await this.getContextId();
     if (!contextId) {
       return { files: [], docs: [] };
@@ -586,3 +706,9 @@ export abstract class AIChatInput extends SignalWatcher(
     };
   }
 }
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'ai-chat-input': AIChatInput;
+  }
+}
diff --git a/packages/frontend/core/src/blocksuite/ai/effects.ts b/packages/frontend/core/src/blocksuite/ai/effects.ts
index 71f2993f62..d1234eeec5 100644
--- a/packages/frontend/core/src/blocksuite/ai/effects.ts
+++ b/packages/frontend/core/src/blocksuite/ai/effects.ts
@@ -23,7 +23,6 @@ import { ActionMindmap } from './chat-panel/actions/mindmap';
 import { ActionSlides } from './chat-panel/actions/slides';
 import { ActionText } from './chat-panel/actions/text';
 import { AILoading } from './chat-panel/ai-loading';
-import { ChatPanelInput } from './chat-panel/chat-panel-input';
 import { ChatPanelMessages } from './chat-panel/chat-panel-messages';
 import { AssistantAvatar } from './chat-panel/content/assistant-avatar';
 import { ChatContentImages } from './chat-panel/content/images';
@@ -40,6 +39,7 @@ import { ChatPanelCollectionChip } from './components/ai-chat-chips/collection-c
 import { ChatPanelDocChip } from './components/ai-chat-chips/doc-chip';
 import { ChatPanelFileChip } from './components/ai-chat-chips/file-chip';
 import { ChatPanelTagChip } from './components/ai-chat-chips/tag-chip';
+import { AIChatInput } from './components/ai-chat-input/ai-chat-input';
 import { effects as componentAiItemEffects } from './components/ai-item';
 import { AIScrollableTextRenderer } from './components/ai-scrollable-text-renderer';
 import { AskAIButton } from './components/ask-ai-button';
@@ -54,7 +54,6 @@ import { AIErrorWrapper } from './messages/error';
 import { AISlidesRenderer } from './messages/slides-renderer';
 import { AIAnswerWrapper } from './messages/wrapper';
 import { registerMiniMindmapBlocks } from './mini-mindmap';
-import { ChatBlockInput } from './peek-view/chat-block-input';
 import { AIChatBlockPeekView } from './peek-view/chat-block-peek-view';
 import { DateTime } from './peek-view/date-time';
 import {
@@ -96,9 +95,9 @@ export function registerAIEffects() {
   customElements.define('action-slides', ActionSlides);
   customElements.define('action-text', ActionText);
   customElements.define('ai-loading', AILoading);
-  customElements.define('chat-panel-input', ChatPanelInput);
   customElements.define('chat-panel-messages', ChatPanelMessages);
   customElements.define('chat-panel', ChatPanel);
+  customElements.define('ai-chat-input', AIChatInput);
   customElements.define('chat-panel-chips', ChatPanelChips);
   customElements.define('chat-panel-add-popover', ChatPanelAddPopover);
   customElements.define(
@@ -113,7 +112,6 @@ export function registerAIEffects() {
   customElements.define('ai-error-wrapper', AIErrorWrapper);
   customElements.define('ai-slides-renderer', AISlidesRenderer);
   customElements.define('ai-answer-wrapper', AIAnswerWrapper);
-  customElements.define('chat-block-input', ChatBlockInput);
   customElements.define('ai-chat-block-peek-view', AIChatBlockPeekView);
   customElements.define('date-time', DateTime);
   customElements.define(
diff --git 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
deleted file mode 100644
index 54497c4e23..0000000000
--- a/packages/frontend/core/src/blocksuite/ai/peek-view/chat-block-input.ts
+++ /dev/null
@@ -1,119 +0,0 @@
-import { property } from 'lit/decorators.js';
-
-import { AIChatInput } from '../components/ai-chat-input';
-import type { ChatMessage } from '../components/ai-chat-messages';
-import { type AIError, AIProvider } from '../provider';
-import { readBlobAsURL } from '../utils/image';
-
-export class ChatBlockInput extends AIChatInput {
-  @property({ attribute: false })
-  accessor getBlockId!: () => string | null | undefined;
-
-  @property({ attribute: false })
-  accessor updateChatBlock!: () => Promise<void>;
-
-  @property({ attribute: false })
-  accessor createChatBlock!: () => Promise<void>;
-
-  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;
-    if (!text) return;
-
-    try {
-      const { doc } = this.host;
-      const promptName = this.getPromptName();
-
-      this.updateContext({
-        images: [],
-        status: 'loading',
-        error: null,
-      });
-
-      const attachments = await Promise.all(
-        images?.map(image => readBlobAsURL(image))
-      );
-
-      const userInfo = await AIProvider.userInfo;
-      this.updateContext({
-        messages: [
-          ...this.chatContextValue.messages,
-          {
-            id: '',
-            content: text,
-            role: 'user',
-            createdAt: new Date().toISOString(),
-            attachments,
-            userId: userInfo?.id,
-            userName: userInfo?.name,
-            avatarUrl: userInfo?.avatarUrl ?? undefined,
-          },
-          {
-            id: '',
-            content: '',
-            role: 'assistant',
-            createdAt: new Date().toISOString(),
-          },
-        ],
-      });
-
-      // must update prompt name after local chat message is updated
-      // otherwise, the unauthorized error can not be rendered properly
-      await this.updatePromptName(promptName);
-
-      const abortController = new AbortController();
-      const stream = AIProvider.actions.chat?.({
-        input: text,
-        sessionId,
-        docId: doc.id,
-        attachments: images,
-        workspaceId: doc.workspace.id,
-        host: this.host,
-        stream: true,
-        signal: abortController.signal,
-        where: 'ai-chat-block',
-        control: 'chat-send',
-      });
-
-      if (stream) {
-        this.updateContext({
-          abortController,
-        });
-
-        for await (const text of stream) {
-          const messages = [...this.chatContextValue.messages];
-          const last = messages[messages.length - 1] as ChatMessage;
-          last.content += text;
-          this.updateContext({ messages, status: 'transmitting' });
-          content += text;
-        }
-
-        this.updateContext({ status: 'success' });
-      }
-    } catch (error) {
-      console.error(error);
-      this.updateContext({ status: 'error', error: error as AIError });
-    } finally {
-      if (content) {
-        const chatBlockExists = !!this.getBlockId();
-        if (!chatBlockExists) {
-          await this.createChatBlock();
-        }
-        // Update new chat block messages if there are contents returned from AI
-        await this.updateChatBlock();
-      }
-
-      this.updateContext({ abortController: null });
-    }
-  };
-}
-
-declare global {
-  interface HTMLElementTagNameMap {
-    'chat-block-input': ChatBlockInput;
-  }
-}
diff --git 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
index e31d76b2a4..973dd630d9 100644
--- 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
@@ -173,8 +173,12 @@ export class AIChatBlockPeekView extends LitElement {
     return this._chatContextId;
   };
 
-  private readonly _getBlockId = () => {
-    return this._chatBlockId;
+  private readonly _onChatSuccess = async () => {
+    if (!this._chatBlockId) {
+      await this.createAIChatBlock();
+    }
+    // Update new chat block messages if there are contents returned from AI
+    await this.updateChatBlockMessages();
   };
 
   /**
@@ -518,8 +522,6 @@ export class AIChatBlockPeekView extends LitElement {
     const latestHistoryMessage = _historyMessages[_historyMessages.length - 1];
     const latestMessageCreatedAt = latestHistoryMessage.createdAt;
     const {
-      updateChatBlockMessages,
-      createAIChatBlock,
       cleanCurrentChatHistories,
       chatContext,
       updateContext,
@@ -541,20 +543,22 @@ export class AIChatBlockPeekView extends LitElement {
           ${this.CurrentMessages(currentChatMessages)}
         </div>
       </div>
-      <chat-block-input
+      <ai-chat-input
         .host=${host}
         .chips=${this.chips}
         .getSessionId=${this._getSessionId}
         .getContextId=${this._getContextId}
-        .getBlockId=${this._getBlockId}
-        .updateChatBlock=${updateChatBlockMessages}
-        .createChatBlock=${createAIChatBlock}
         .cleanupHistories=${cleanCurrentChatHistories}
         .chatContextValue=${chatContext}
         .updateContext=${updateContext}
         .networkSearchConfig=${networkSearchConfig}
         .docDisplayConfig=${this.docDisplayConfig}
-      ></chat-block-input>
+        .onChatSuccess=${this._onChatSuccess}
+        .trackOptions=${{
+          where: 'ai-chat-block',
+          control: 'chat-send',
+        }}
+      ></ai-chat-input>
       <div class="peek-view-footer">
         ${InformationIcon()}
         <div>AI outputs can be misleading or wrong</div>
diff --git a/tests/affine-cloud-copilot/e2e/chat-with/network.spec.ts b/tests/affine-cloud-copilot/e2e/chat-with/network.spec.ts
index a4fd86c650..e25c3fe648 100644
--- a/tests/affine-cloud-copilot/e2e/chat-with/network.spec.ts
+++ b/tests/affine-cloud-copilot/e2e/chat-with/network.spec.ts
@@ -96,7 +96,7 @@ test.describe('AIChatWith/Network', () => {
       buffer: buffer,
     });
 
-    await page.waitForSelector('chat-panel-input img');
+    await page.waitForSelector('ai-chat-input img');
     const isNetworkSearchEnabled =
       await utils.chatPanel.isNetworkSearchEnabled(page);
     await expect(isNetworkSearchEnabled).toBe(false);
diff --git a/tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts b/tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts
index 122f66dcab..e550c54320 100644
--- a/tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts
+++ b/tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts
@@ -244,7 +244,7 @@ export class ChatPanelUtils {
     const fileChooser = await fileChooserPromise;
     await fileChooser.setFiles(images);
 
-    await page.waitForSelector('chat-panel-input img');
+    await page.waitForSelector('ai-chat-input img');
     await this.makeChat(page, text);
   }