feat(core): support objects render in ai chat block (#12906)

Close [AI-195](https://linear.app/affine-design/issue/AI-195)


<img width="1219" alt="截屏2025-06-24 10 55 18"
src="https://github.com/user-attachments/assets/8c54ad41-2bbf-443a-a41a-9cea9aede7b4"
/>


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Introduced a new component to render various AI stream objects within
chat messages, supporting richer content types like reasoning blocks and
tool calls.

- **Improvements**
- Enhanced chat message rendering by consolidating message properties
into a single object and simplifying component interfaces.
- Improved scrolling behavior with smooth, throttled scrolling and
better synchronization with chat status.
- Updated chat history and streaming updates to use immutable data
handling for more reliable message updates.
- Refined rendering keys and content merging for more accurate display
of streamed AI responses.
- Improved handling and display of tool results and web search/crawl
outputs.

- **Bug Fixes**
- Ensured consistent presence of attachments and streaming objects in
chat message data.
- Enhanced parsing and merging of streamed AI responses for better
message accuracy.

- **Refactor**
- Streamlined imports, component registrations, and internal message
handling for maintainability.
- Updated type definitions to support new optional message properties
and improve compatibility.

- **Style**
- Adjusted layout by moving history clear controls for a more intuitive
chat interface.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/backend/server/src/models/copilot-session.ts

@@ -194,6 +194,7 @@ export class CopilotSessionModel extends BaseModel {
           id: true,
           role: true,
           content: true,
+          streamObjects: true,
           attachments: true,
           params: true,
           createdAt: true,

packages/frontend/core/src/blocksuite/ai/_common/chat-actions-handle.ts

@@ -91,7 +91,7 @@ export function constructUserInfoWithMessages(
   userInfo: AIUserInfo | null
 ) {
   return messages.map(message => {
-    const { role, id, content, createdAt } = message;
+    const { role, streamObjects } = message;
     const isUser = role === 'user';
     const userInfoProps = isUser
       ? {
@@ -101,12 +101,10 @@ export function constructUserInfoWithMessages(
         }
       : {};
     return {
-      id,
-      role,
-      content,
-      createdAt,
-      attachments: [],
+      ...message,
       ...userInfoProps,
+      attachments: [],
+      streamObjects: streamObjects || [],
     };
   });
 }
@@ -117,11 +115,11 @@ export async function constructRootChatBlockMessages(
 ) {
   // Convert chat messages to AI chat block messages
   const userInfo = await AIProvider.userInfo;
-  const forkMessages = await queryHistoryMessages(
+  const forkMessages = (await queryHistoryMessages(
     doc.workspace.id,
     doc.id,
     forkSessionId
-  );
+  )) as ChatMessage[];
   return constructUserInfoWithMessages(forkMessages, userInfo);
 }
 

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

@@ -9,6 +9,7 @@ import type {
   CopilotSessionType,
   getCopilotHistoriesQuery,
   RequestOptions,
+  StreamObject,
 } from '@affine/graphql';
 import type { EditorHost } from '@blocksuite/affine/std';
 import type { GfxModel } from '@blocksuite/affine/std/gfx';
@@ -359,7 +360,8 @@ declare global {
         content: string;
         createdAt: string;
         role: MessageRole;
-        attachments?: string[];
+        attachments?: string[] | null;
+        streamObjects?: StreamObject[] | null;
       }[];
     }
 

packages/frontend/core/src/blocksuite/ai/blocks/ai-chat-block/components/ai-chat-messages.ts

@@ -1,14 +1,13 @@
 import type { TextRendererOptions } from '@affine/core/blocksuite/ai/components/text-renderer';
-import type { AffineAIPanelState } from '@affine/core/blocksuite/ai/widgets/ai-panel/type';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { css, html, LitElement } from 'lit';
 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,
+import {
+  type ChatMessage,
+  type StreamObject,
 } from '../../../components/ai-chat-messages';
 import { UserInfoTemplate } from './user-info';
 
@@ -36,16 +35,14 @@ export class AIChatMessage extends LitElement {
 
   override render() {
     const {
-      host,
-      textRendererOptions,
-      state,
       content,
       attachments,
-      messageRole,
-      userId,
       userName,
+      userId,
       avatarUrl,
-    } = this;
+      role,
+      streamObjects,
+    } = this.message;
     const withAttachments = !!attachments && attachments.length > 0;
 
     const messageClasses = classMap({
@@ -54,48 +51,52 @@ export class AIChatMessage extends LitElement {
 
     return html`
       <div class="ai-chat-message">
-        ${UserInfoTemplate({ userId, userName, avatarUrl }, messageRole)}
+        ${UserInfoTemplate({ userId, userName, avatarUrl }, role)}
         <div class="ai-chat-content">
           <chat-images .attachments=${attachments}></chat-images>
           <div class=${messageClasses}>
-            <text-renderer
-              .host=${host}
-              .answer=${content}
-              .options=${textRendererOptions}
-              .state=${state}
-            ></text-renderer>
+            ${streamObjects?.length
+              ? this.renderStreamObjects(streamObjects)
+              : this.renderRichText(content)}
           </div>
         </div>
       </div>
     `;
   }
 
-  @property({ attribute: false })
-  accessor attachments: string[] | undefined = undefined;
+  private renderStreamObjects(answer: StreamObject[]) {
+    return html`<chat-content-stream-objects
+      .answer=${answer}
+      .host=${this.host}
+      .state=${this.state}
+      .extensions=${this.textRendererOptions.extensions}
+      .affineFeatureFlagService=${this.textRendererOptions
+        .affineFeatureFlagService}
+    ></chat-content-stream-objects>`;
+  }
+
+  private renderRichText(text: string) {
+    return html`<chat-content-rich-text
+      .host=${this.host}
+      .text=${text}
+      .state=${this.state}
+      .extensions=${this.textRendererOptions.extensions}
+      .affineFeatureFlagService=${this.textRendererOptions
+        .affineFeatureFlagService}
+    ></chat-content-rich-text>`;
+  }
 
   @property({ attribute: false })
-  accessor content: string = '';
+  accessor message!: ChatMessage;
 
   @property({ attribute: false })
   accessor host!: EditorHost;
 
   @property({ attribute: false })
-  accessor messageRole: MessageRole | undefined = undefined;
-
-  @property({ attribute: false })
-  accessor state: AffineAIPanelState = 'finished';
+  accessor state: 'finished' | 'generating' = 'finished';
 
   @property({ attribute: false })
   accessor textRendererOptions: TextRendererOptions = {};
-
-  @property({ attribute: false })
-  accessor userId: string | undefined = undefined;
-
-  @property({ attribute: false })
-  accessor userName: string | undefined = undefined;
-
-  @property({ attribute: false })
-  accessor avatarUrl: string | undefined = undefined;
 }
 
 export class AIChatMessages extends LitElement {
@@ -121,18 +122,11 @@ export class AIChatMessages extends LitElement {
         this.messages,
         message => message.id || message.createdAt,
         message => {
-          const { attachments, role, content, userId, userName, avatarUrl } =
-            message;
           return html`
             <ai-chat-message
               .host=${this.host}
               .textRendererOptions=${this.textRendererOptions}
-              .content=${content}
-              .attachments=${attachments}
-              .messageRole=${role}
-              .userId=${userId}
-              .userName=${userName}
-              .avatarUrl=${avatarUrl}
+              .message=${message}
             ></ai-chat-message>
           `;
         }

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

@@ -1,6 +1,3 @@
-import './action-wrapper';
-import '../content/images';
-
 import { WithDisposable } from '@blocksuite/affine/global/lit';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';

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

@@ -19,7 +19,6 @@ import type {
   AIReasoningConfig,
 } from '../components/ai-chat-input';
 import {
-  type ChatMessage,
   isChatAction,
   isChatMessage,
   StreamObjectSchema,
@@ -418,22 +417,27 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
       });
 
       for await (const text of stream) {
-        const messages = [...this.chatContextValue.messages];
-        const last = messages[messages.length - 1] as ChatMessage;
-        try {
-          const parsed = StreamObjectSchema.safeParse(JSON.parse(text));
-          if (parsed.success) {
-            last.streamObjects = mergeStreamObjects([
+        const messages = this.chatContextValue.messages.slice(0);
+        const last = messages.at(-1);
+        if (last && isChatMessage(last)) {
+          try {
+            const parsed = StreamObjectSchema.parse(JSON.parse(text));
+            const streamObjects = mergeStreamObjects([
               ...(last.streamObjects ?? []),
-              parsed.data,
+              parsed,
             ]);
-          } else {
-            last.content += text;
+            messages[messages.length - 1] = {
+              ...last,
+              streamObjects,
+            };
+          } catch {
+            messages[messages.length - 1] = {
+              ...last,
+              content: last.content + text,
+            };
           }
-        } catch {
-          last.content += text;
+          this.updateContext({ messages, status: 'transmitting' });
         }
-        this.updateContext({ messages, status: 'transmitting' });
       }
 
       this.updateContext({ status: 'success' });

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

@@ -24,7 +24,11 @@ import type {
   AINetworkSearchConfig,
   AIReasoningConfig,
 } from '../components/ai-chat-input';
-import { type HistoryMessage } from '../components/ai-chat-messages';
+import {
+  type ChatAction,
+  type ChatMessage,
+  type HistoryMessage,
+} from '../components/ai-chat-messages';
 import { createPlaygroundModal } from '../components/playground/modal';
 import { AIProvider } from '../provider';
 import { extractSelectedContent } from '../utils/extract';
@@ -146,12 +150,14 @@ export class ChatPanel extends SignalWatcher(
       return;
     }
 
-    const messages: HistoryMessage[] = actions ? [...actions] : [];
+    const chatActions = (actions || []) as ChatAction[];
+    const messages: HistoryMessage[] = chatActions;
 
     const sessionId = await this._getSessionId();
     const history = histories?.find(history => history.sessionId === sessionId);
     if (history) {
-      messages.push(...history.messages);
+      const chatMessages = (history.messages || []) as ChatMessage[];
+      messages.push(...chatMessages);
     }
 
     this.chatContextValue = {

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

@@ -1,5 +1,3 @@
-import '../content/assistant-avatar';
-
 import { WithDisposable } from '@blocksuite/affine/global/lit';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';

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

@@ -1,9 +1,5 @@
-import '../content/assistant-avatar';
-import '../content/rich-text';
-
 import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import { WithDisposable } from '@blocksuite/affine/global/lit';
-import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
 import { isInsidePageEditor } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
@@ -33,20 +29,6 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
       font-size: var(--affine-font-xs);
       font-weight: 400;
     }
-
-    .reasoning-wrapper {
-      padding: 16px 20px;
-      margin: 8px 0;
-      border-radius: 8px;
-      background-color: rgba(0, 0, 0, 0.05);
-    }
-
-    .tool-wrapper {
-      padding: 12px;
-      margin: 8px 0;
-      border-radius: 8px;
-      border: 0.5px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
-    }
   `;
 
   @property({ attribute: false })
@@ -82,6 +64,15 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
   @property({ attribute: false })
   accessor panelWidth!: Signal<number | undefined>;
 
+  get state() {
+    const { isLast, status } = this;
+    return isLast
+      ? status !== 'loading' && status !== 'transmitting'
+        ? 'finished'
+        : 'generating'
+      : 'finished';
+  }
+
   renderHeader() {
     const isWithDocs =
       'content' in this.item &&
@@ -122,102 +113,21 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
   }
 
   private renderStreamObjects(answer: StreamObject[]) {
-    return html`<div>
-      ${answer.map(data => {
-        switch (data.type) {
-          case 'text-delta':
-            return this.renderRichText(data.textDelta);
-          case 'reasoning':
-            return html`
-              <div class="reasoning-wrapper">
-                ${this.renderRichText(data.textDelta)}
-              </div>
-            `;
-          case 'tool-call':
-            return this.renderToolCall(data);
-          case 'tool-result':
-            return this.renderToolResult(data);
-          default:
-            return nothing;
-        }
-      })}
-    </div>`;
-  }
-
-  private renderToolCall(streamObject: StreamObject) {
-    if (streamObject.type !== 'tool-call') {
-      return nothing;
-    }
-
-    switch (streamObject.toolName) {
-      case 'web_crawl_exa':
-        return html`
-          <web-crawl-tool
-            .data=${streamObject}
-            .host=${this.host}
-            .width=${this.panelWidth}
-          ></web-crawl-tool>
-        `;
-      case 'web_search_exa':
-        return html`
-          <web-search-tool
-            .data=${streamObject}
-            .host=${this.host}
-            .width=${this.panelWidth}
-          ></web-search-tool>
-        `;
-      default:
-        return html`
-          <div class="tool-wrapper">
-            ${streamObject.toolName} tool calling...
-          </div>
-        `;
-    }
-  }
-
-  private renderToolResult(streamObject: StreamObject) {
-    if (streamObject.type !== 'tool-result') {
-      return nothing;
-    }
-
-    switch (streamObject.toolName) {
-      case 'web_crawl_exa':
-        return html`
-          <web-crawl-tool
-            .data=${streamObject}
-            .host=${this.host}
-            .width=${this.panelWidth}
-          ></web-crawl-tool>
-        `;
-      case 'web_search_exa':
-        return html`
-          <web-search-tool
-            .data=${streamObject}
-            .host=${this.host}
-            .width=${this.panelWidth}
-          ></web-search-tool>
-        `;
-      default:
-        return html`
-          <div class="tool-wrapper">
-            ${streamObject.toolName} tool result...
-          </div>
-        `;
-    }
+    return html`<chat-content-stream-objects
+      .answer=${answer}
+      .host=${this.host}
+      .state=${this.state}
+      .width=${this.panelWidth}
+      .extensions=${this.extensions}
+      .affineFeatureFlagService=${this.affineFeatureFlagService}
+    ></chat-content-stream-objects>`;
   }
 
   private renderRichText(text: string) {
-    const { host, isLast, status } = this;
-    const state = isLast
-      ? status !== 'loading' && status !== 'transmitting'
-        ? 'finished'
-        : 'generating'
-      : 'finished';
-
     return html`<chat-content-rich-text
-      .host=${host}
+      .host=${this.host}
       .text=${text}
-      .state=${state}
+      .state=${this.state}
       .extensions=${this.extensions}
       .affineFeatureFlagService=${this.affineFeatureFlagService}
     ></chat-content-rich-text>`;

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

@@ -1,6 +1,3 @@
-import '../content/images';
-import '../content/pure-text';
-
 import { WithDisposable } from '@blocksuite/affine/global/lit';
 import { ShadowlessElement } from '@blocksuite/affine/std';
 import { css, html, nothing } from 'lit';

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

@@ -19,7 +19,11 @@ import { readBlobAsURL } from '../../utils/image';
 import { mergeStreamObjects } from '../../utils/stream-objects';
 import type { ChatChip, DocDisplayConfig } from '../ai-chat-chips/type';
 import { isDocChip } from '../ai-chat-chips/utils';
-import { type ChatMessage, StreamObjectSchema } from '../ai-chat-messages';
+import {
+  type ChatMessage,
+  isChatMessage,
+  StreamObjectSchema,
+} from '../ai-chat-messages';
 import { MAX_IMAGE_COUNT } from './const';
 import type {
   AIChatInputContext,
@@ -609,22 +613,27 @@ export class AIChatInput extends SignalWatcher(
       });
 
       for await (const text of stream) {
-        const messages = [...this.chatContextValue.messages];
-        const last = messages[messages.length - 1] as ChatMessage;
-        try {
-          const parsed = StreamObjectSchema.safeParse(JSON.parse(text));
-          if (parsed.success) {
-            last.streamObjects = mergeStreamObjects([
+        const messages = this.chatContextValue.messages.slice(0);
+        const last = messages.at(-1);
+        if (last && isChatMessage(last)) {
+          try {
+            const parsed = StreamObjectSchema.parse(JSON.parse(text));
+            const streamObjects = mergeStreamObjects([
               ...(last.streamObjects ?? []),
-              parsed.data,
+              parsed,
             ]);
-          } else {
-            last.content += text;
+            messages[messages.length - 1] = {
+              ...last,
+              streamObjects,
+            };
+          } catch {
+            messages[messages.length - 1] = {
+              ...last,
+              content: last.content + text,
+            };
           }
-        } catch {
-          last.content += text;
+          this.updateContext({ messages, status: 'transmitting' });
         }
-        this.updateContext({ messages, status: 'transmitting' });
       }
 
       this.updateContext({ status: 'success' });

packages/frontend/core/src/blocksuite/ai/components/ai-message-content/assistant-avatar.ts

@@ -4,7 +4,7 @@ import { AiIcon } from '@blocksuite/icons/lit';
 import { css, html } from 'lit';
 import { property } from 'lit/decorators.js';
 
-import type { ChatStatus } from '../../components/ai-chat-messages';
+import type { ChatStatus } from '../ai-chat-messages';
 
 const AffineAvatarIcon = AiIcon({
   width: '20px',

packages/frontend/core/src/blocksuite/ai/components/ai-message-content/stream-objects.ts

@@ -0,0 +1,145 @@
+import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
+import { WithDisposable } from '@blocksuite/affine/global/lit';
+import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
+import type { EditorHost } from '@blocksuite/affine/std';
+import { ShadowlessElement } from '@blocksuite/affine/std';
+import type { ExtensionType } from '@blocksuite/affine/store';
+import type { Signal } from '@preact/signals-core';
+import { css, html, nothing } from 'lit';
+import { property } from 'lit/decorators.js';
+
+import type { AffineAIPanelState } from '../../widgets/ai-panel/type';
+import type { StreamObject } from '../ai-chat-messages';
+
+export class ChatContentStreamObjects extends WithDisposable(
+  ShadowlessElement
+) {
+  static override styles = css`
+    .reasoning-wrapper {
+      padding: 16px 20px;
+      margin: 8px 0;
+      border-radius: 8px;
+      background-color: rgba(0, 0, 0, 0.05);
+    }
+
+    .tool-wrapper {
+      padding: 12px;
+      margin: 8px 0;
+      border-radius: 8px;
+      border: 0.5px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
+    }
+  `;
+
+  @property({ attribute: false })
+  accessor answer!: StreamObject[];
+
+  @property({ attribute: false })
+  accessor host!: EditorHost;
+
+  @property({ attribute: false })
+  accessor state: AffineAIPanelState = 'finished';
+
+  @property({ attribute: false })
+  accessor width: Signal<number | undefined> | undefined;
+
+  @property({ attribute: false })
+  accessor extensions!: ExtensionType[];
+
+  @property({ attribute: false })
+  accessor affineFeatureFlagService!: FeatureFlagService;
+
+  private renderToolCall(streamObject: StreamObject) {
+    if (streamObject.type !== 'tool-call') {
+      return nothing;
+    }
+
+    switch (streamObject.toolName) {
+      case 'web_crawl_exa':
+        return html`
+          <web-crawl-tool
+            .data=${streamObject}
+            .host=${this.host}
+            .width=${this.width}
+          ></web-crawl-tool>
+        `;
+      case 'web_search_exa':
+        return html`
+          <web-search-tool
+            .data=${streamObject}
+            .host=${this.host}
+            .width=${this.width}
+          ></web-search-tool>
+        `;
+      default:
+        return html`
+          <div class="tool-wrapper">
+            ${streamObject.toolName} tool calling...
+          </div>
+        `;
+    }
+  }
+
+  private renderToolResult(streamObject: StreamObject) {
+    if (streamObject.type !== 'tool-result') {
+      return nothing;
+    }
+
+    switch (streamObject.toolName) {
+      case 'web_crawl_exa':
+        return html`
+          <web-crawl-tool
+            .data=${streamObject}
+            .host=${this.host}
+            .width=${this.width}
+          ></web-crawl-tool>
+        `;
+      case 'web_search_exa':
+        return html`
+          <web-search-tool
+            .data=${streamObject}
+            .host=${this.host}
+            .width=${this.width}
+          ></web-search-tool>
+        `;
+      default:
+        return html`
+          <div class="tool-wrapper">
+            ${streamObject.toolName} tool result...
+          </div>
+        `;
+    }
+  }
+
+  private renderRichText(text: string) {
+    return html`<chat-content-rich-text
+      .host=${this.host}
+      .text=${text}
+      .state=${this.state}
+      .extensions=${this.extensions}
+      .affineFeatureFlagService=${this.affineFeatureFlagService}
+    ></chat-content-rich-text>`;
+  }
+
+  protected override render() {
+    return html`<div>
+      ${this.answer.map(data => {
+        switch (data.type) {
+          case 'text-delta':
+            return this.renderRichText(data.textDelta);
+          case 'reasoning':
+            return html`
+              <div class="reasoning-wrapper">
+                ${this.renderRichText(data.textDelta)}
+              </div>
+            `;
+          case 'tool-call':
+            return this.renderToolCall(data);
+          case 'tool-result':
+            return this.renderToolResult(data);
+          default:
+            return nothing;
+        }
+      })}
+    </div>`;
+  }
+}

packages/frontend/core/src/blocksuite/ai/components/ai-tools/tool-result-card.ts

@@ -3,7 +3,7 @@ import { ImageProxyService } from '@blocksuite/affine/shared/adapters';
 import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
 import { type EditorHost, ShadowlessElement } from '@blocksuite/affine/std';
 import { ToggleDownIcon } from '@blocksuite/icons/lit';
-import { type Signal, signal } from '@preact/signals-core';
+import { type Signal } from '@preact/signals-core';
 import { css, html, nothing, type TemplateResult } from 'lit';
 import { property, state } from 'lit/decorators.js';
 
@@ -170,7 +170,7 @@ export class ToolResultCard extends SignalWatcher(
   accessor results!: ToolResult[];
 
   @property({ attribute: false })
-  accessor width: Signal<number | undefined> = signal(undefined);
+  accessor width: Signal<number | undefined> | undefined;
 
   @state()
   private accessor isCollapsed = true;
@@ -209,7 +209,10 @@ export class ToolResultCard extends SignalWatcher(
       return nothing;
     }
 
-    const maxIcons = Number(this.width.value) <= 400 ? 1 : 3;
+    let maxIcons = 3;
+    if (this.width && this.width.value !== undefined) {
+      maxIcons = this.width.value <= 400 ? 1 : 3;
+    }
     const visibleIcons = this.footerIcons.slice(0, maxIcons);
 
     return html`

packages/frontend/core/src/blocksuite/ai/components/ai-tools/web-crawl.ts

@@ -35,7 +35,7 @@ export class WebCrawlTool extends WithDisposable(ShadowlessElement) {
   accessor host!: EditorHost;
 
   @property({ attribute: false })
-  accessor width!: Signal<number | undefined>;
+  accessor width: Signal<number | undefined> | undefined;
 
   renderToolCall() {
     return html`

packages/frontend/core/src/blocksuite/ai/components/ai-tools/web-search.ts

@@ -35,7 +35,7 @@ export class WebSearchTool extends WithDisposable(ShadowlessElement) {
   accessor host!: EditorHost;
 
   @property({ attribute: false })
-  accessor width!: Signal<number | undefined>;
+  accessor width: Signal<number | undefined> | undefined;
 
   renderToolCall() {
     return html`

packages/frontend/core/src/blocksuite/ai/components/playground/chat.ts

@@ -22,7 +22,11 @@ import type {
   AINetworkSearchConfig,
   AIReasoningConfig,
 } from '../ai-chat-input';
-import { type HistoryMessage } from '../ai-chat-messages';
+import {
+  type ChatAction,
+  type ChatMessage,
+  type HistoryMessage,
+} from '../ai-chat-messages';
 
 const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
   quote: '',
@@ -204,12 +208,14 @@ export class PlaygroundChat extends SignalWatcher(
       return;
     }
 
-    const messages: HistoryMessage[] = actions ? [...actions] : [];
+    const chatActions = (actions || []) as ChatAction[];
+    const messages: HistoryMessage[] = chatActions;
 
     const sessionId = await this._getSessionId();
     const history = histories?.find(history => history.sessionId === sessionId);
     if (history) {
-      messages.push(...history.messages);
+      const chatMessages = (history.messages || []) as ChatMessage[];
+      messages.push(...chatMessages);
     }
 
     this.chatContextValue = {

packages/frontend/core/src/blocksuite/ai/effects.ts

@@ -21,17 +21,9 @@ import { ActionSlides } from './chat-panel/actions/slides';
 import { ActionText } from './chat-panel/actions/text';
 import { AILoading } from './chat-panel/ai-loading';
 import { ChatPanelMessages } from './chat-panel/chat-panel-messages';
-import { AssistantAvatar } from './chat-panel/content/assistant-avatar';
-import { ChatContentImages } from './chat-panel/content/images';
-import { ChatContentPureText } from './chat-panel/content/pure-text';
-import { ChatContentRichText } from './chat-panel/content/rich-text';
 import { ChatMessageAction } from './chat-panel/message/action';
 import { ChatMessageAssistant } from './chat-panel/message/assistant';
 import { ChatMessageUser } from './chat-panel/message/user';
-import { ToolCallCard } from './chat-panel/tools/tool-call-card';
-import { ToolResultCard } from './chat-panel/tools/tool-result-card';
-import { WebCrawlTool } from './chat-panel/tools/web-crawl';
-import { WebSearchTool } from './chat-panel/tools/web-search';
 import { ChatPanelAddPopover } from './components/ai-chat-chips/add-popover';
 import { ChatPanelCandidatesPopover } from './components/ai-chat-chips/candidates-popover';
 import { ChatPanelChips } from './components/ai-chat-chips/chat-panel-chips';
@@ -46,7 +38,16 @@ import { AIChatEmbeddingStatusTooltip } from './components/ai-chat-input/embeddi
 import { ChatInputPreference } from './components/ai-chat-input/preference-popup';
 import { AIHistoryClear } from './components/ai-history-clear';
 import { effects as componentAiItemEffects } from './components/ai-item';
+import { AssistantAvatar } from './components/ai-message-content/assistant-avatar';
+import { ChatContentImages } from './components/ai-message-content/images';
+import { ChatContentPureText } from './components/ai-message-content/pure-text';
+import { ChatContentRichText } from './components/ai-message-content/rich-text';
+import { ChatContentStreamObjects } from './components/ai-message-content/stream-objects';
 import { AIScrollableTextRenderer } from './components/ai-scrollable-text-renderer';
+import { ToolCallCard } from './components/ai-tools/tool-call-card';
+import { ToolResultCard } from './components/ai-tools/tool-result-card';
+import { WebCrawlTool } from './components/ai-tools/web-crawl';
+import { WebSearchTool } from './components/ai-tools/web-search';
 import { AskAIButton } from './components/ask-ai-button';
 import { AskAIIcon } from './components/ask-ai-icon';
 import { AskAIPanel } from './components/ask-ai-panel';
@@ -156,6 +157,10 @@ export function registerAIEffects() {
   customElements.define('chat-content-images', ChatContentImages);
   customElements.define('chat-content-pure-text', ChatContentPureText);
   customElements.define('chat-content-rich-text', ChatContentRichText);
+  customElements.define(
+    'chat-content-stream-objects',
+    ChatContentStreamObjects
+  );
   customElements.define('chat-message-action', ChatMessageAction);
   customElements.define('chat-message-assistant', ChatMessageAssistant);
   customElements.define('chat-message-user', ChatMessageUser);

packages/frontend/core/src/blocksuite/ai/peek-view/chat-block-peek-view.ts

@@ -14,10 +14,11 @@ import {
 import type { Signal } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { signal } from '@preact/signals-core';
-import { html, LitElement, nothing } from 'lit';
+import { html, LitElement, nothing, type PropertyValues } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { classMap } from 'lit/directives/class-map.js';
 import { repeat } from 'lit/directives/repeat.js';
+import { throttle } from 'lodash-es';
 
 import {
   ChatBlockPeekViewActions,
@@ -36,12 +37,16 @@ import type {
 import type { ChatMessage } from '../components/ai-chat-messages';
 import {
   ChatMessagesSchema,
+  isChatMessage,
   StreamObjectSchema,
 } from '../components/ai-chat-messages';
 import type { TextRendererOptions } from '../components/text-renderer';
 import { AIChatErrorRenderer } from '../messages/error';
 import { type AIError, AIProvider } from '../provider';
-import { mergeStreamObjects } from '../utils/stream-objects';
+import {
+  mergeStreamContent,
+  mergeStreamObjects,
+} from '../utils/stream-objects';
 import { PeekViewStyles } from './styles';
 import type { ChatContext } from './types';
 import { calcChildBound } from './utils';
@@ -115,11 +120,11 @@ export class AIChatBlockPeekView extends LitElement {
     forkSessionId: string
   ) => {
     const currentUserInfo = await AIProvider.userInfo;
-    const forkMessages = await queryHistoryMessages(
+    const forkMessages = (await queryHistoryMessages(
       rootWorkspaceId,
       rootDocId,
       forkSessionId
-    );
+    )) as ChatMessage[];
     const forkLength = forkMessages.length;
     const historyLength = this._historyMessages.length;
 
@@ -332,6 +337,18 @@ export class AIChatBlockPeekView extends LitElement {
     this._resetContext();
   };
 
+  private readonly _scrollToEnd = () => {
+    requestAnimationFrame(() => {
+      if (!this._chatMessagesContainer) return;
+      this._chatMessagesContainer.scrollTo({
+        top: this._chatMessagesContainer.scrollHeight,
+        behavior: 'smooth',
+      });
+    });
+  };
+
+  private readonly _throttledScrollToEnd = throttle(this._scrollToEnd, 600);
+
   /**
    * Retry the last chat message
    */
@@ -372,22 +389,27 @@ export class AIChatBlockPeekView extends LitElement {
       });
 
       for await (const text of stream) {
-        const messages = [...this.chatContext.messages];
-        const last = messages[messages.length - 1] as ChatMessage;
-        try {
-          const parsed = StreamObjectSchema.safeParse(JSON.parse(text));
-          if (parsed.success) {
-            last.streamObjects = mergeStreamObjects([
+        const messages = this.chatContext.messages.slice(0);
+        const last = messages.at(-1);
+        if (last && isChatMessage(last)) {
+          try {
+            const parsed = StreamObjectSchema.parse(JSON.parse(text));
+            const streamObjects = mergeStreamObjects([
               ...(last.streamObjects ?? []),
-              parsed.data,
+              parsed,
             ]);
-          } else {
-            last.content += text;
+            messages[messages.length - 1] = {
+              ...last,
+              streamObjects,
+            };
+          } catch {
+            messages[messages.length - 1] = {
+              ...last,
+              content: last.content + text,
+            };
           }
-        } catch {
-          last.content += text;
+          this.updateContext({ messages, status: 'transmitting' });
         }
-        this.updateContext({ messages, status: 'transmitting' });
       }
 
       this.updateContext({ status: 'success' });
@@ -410,7 +432,7 @@ export class AIChatBlockPeekView extends LitElement {
 
     return html`${repeat(
       currentMessages,
-      message => message.id || message.createdAt,
+      (_, index) => index,
       (message, idx) => {
         const { status, error } = this.chatContext;
         const isAssistantMessage = message.role === 'assistant';
@@ -424,26 +446,24 @@ export class AIChatBlockPeekView extends LitElement {
         const isNotReady = status === 'transmitting' || status === 'loading';
         const shouldRenderCopyMore =
           isAssistantMessage && !(isLastReply && isNotReady);
-        const shouldRenderActions =
-          isLastReply && !!message.content && !isNotReady;
+        const markdown = message.streamObjects?.length
+          ? mergeStreamContent(message.streamObjects)
+          : message.content;
+        const shouldRenderActions = isLastReply && !!markdown && !isNotReady;
 
         const messageClasses = classMap({
           'assistant-message-container': isAssistantMessage,
         });
 
-        const { attachments, role, content, userId, userName, avatarUrl } =
-          message;
+        if (status === 'loading' && isLastReply) {
+          return html`<ai-loading></ai-loading>`;
+        }
 
         return html`<div class=${messageClasses}>
           <ai-chat-message
             .host=${host}
             .state=${messageState}
-            .content=${content}
-            .attachments=${attachments}
-            .messageRole=${role}
-            .userId=${userId}
-            .userName=${userName}
-            .avatarUrl=${avatarUrl}
+            .message=${message}
             .textRendererOptions=${this._textRendererOptions}
           ></ai-chat-message>
           ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
@@ -451,7 +471,7 @@ export class AIChatBlockPeekView extends LitElement {
             ? html` <chat-copy-more
                 .host=${host}
                 .actions=${actions}
-                .content=${message.content}
+                .content=${markdown}
                 .isLast=${isLastReply}
                 .getSessionId=${this._getSessionId}
                 .messageId=${message.id ?? undefined}
@@ -462,7 +482,7 @@ export class AIChatBlockPeekView extends LitElement {
             ? html`<chat-action-list
                 .host=${host}
                 .actions=${actions}
-                .content=${message.content}
+                .content=${markdown}
                 .getSessionId=${this._getSessionId}
                 .messageId=${message.id ?? undefined}
                 .layoutDirection=${'horizontal'}
@@ -502,13 +522,25 @@ export class AIChatBlockPeekView extends LitElement {
   }
 
   override firstUpdated() {
-    // first time render, scroll ai-chat-messages-container to bottom
-    requestAnimationFrame(() => {
-      if (this._chatMessagesContainer) {
-        this._chatMessagesContainer.scrollTop =
-          this._chatMessagesContainer.scrollHeight;
-      }
-    });
+    this._scrollToEnd();
+  }
+
+  protected override updated(changedProperties: PropertyValues) {
+    if (
+      changedProperties.has('chatContext') &&
+      (this.chatContext.status === 'loading' ||
+        this.chatContext.status === 'error' ||
+        this.chatContext.status === 'success')
+    ) {
+      setTimeout(this._scrollToEnd, 500);
+    }
+
+    if (
+      changedProperties.has('chatContext') &&
+      this.chatContext.status === 'transmitting'
+    ) {
+      this._throttledScrollToEnd();
+    }
   }
 
   override render() {
@@ -529,6 +561,15 @@ export class AIChatBlockPeekView extends LitElement {
     const { messages: currentChatMessages } = chatContext;
 
     return html`<div class="ai-chat-block-peek-view-container">
+      <div class="history-clear-container">
+        <ai-history-clear
+          .host=${this.host}
+          .doc=${this.host.store}
+          .getSessionId=${this._getSessionId}
+          .onHistoryCleared=${this._onHistoryCleared}
+          .chatContextValue=${chatContext}
+        ></ai-history-clear>
+      </div>
       <div class="ai-chat-messages-container">
         <ai-chat-messages
           .host=${host}
@@ -539,15 +580,6 @@ export class AIChatBlockPeekView extends LitElement {
         <div class="new-chat-messages-container">
           ${this.CurrentMessages(currentChatMessages)}
         </div>
-        <div class="history-clear-container">
-          <ai-history-clear
-            .host=${this.host}
-            .doc=${this.host.store}
-            .getSessionId=${this._getSessionId}
-            .onHistoryCleared=${this._onHistoryCleared}
-            .chatContextValue=${chatContext}
-          ></ai-history-clear>
-        </div>
       </div>
       <ai-chat-composer
         .host=${host}

packages/frontend/core/src/blocksuite/ai/utils/stream-objects.ts

@@ -7,7 +7,10 @@ export function mergeStreamObjects(chunks: StreamObject[] = []) {
       case 'reasoning':
       case 'text-delta': {
         if (prev && prev.type === curr.type) {
-          prev.textDelta += curr.textDelta;
+          acc[acc.length - 1] = {
+            ...prev,
+            textDelta: prev.textDelta + curr.textDelta,
+          };
         } else {
           acc.push(curr);
         }