feat: support chatting in center peek (#7601)

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

@@ -0,0 +1,506 @@
+import { ChatHistoryOrder } from '@affine/graphql';
+import type {
+  BlockSelection,
+  EditorHost,
+  TextSelection,
+} from '@blocksuite/block-std';
+import type {
+  DocMode,
+  EdgelessRootService,
+  ImageSelection,
+  PageRootService,
+} from '@blocksuite/blocks';
+import {
+  BlocksUtils,
+  getElementsBound,
+  NoteDisplayMode,
+} from '@blocksuite/blocks';
+import { Bound, type SerializedXYWH } from '@blocksuite/global/utils';
+import { type ChatMessage } from '@blocksuite/presets';
+import type { Doc } from '@blocksuite/store';
+import type { TemplateResult } from 'lit';
+
+import { AIProvider, type AIUserInfo } from '../provider';
+import { reportResponse } from '../utils/action-reporter';
+import { insertBelow, replace } from '../utils/editor-actions';
+import { insertFromMarkdown } from '../utils/markdown-utils';
+import { BlockIcon, CreateIcon, InsertBelowIcon, ReplaceIcon } from './icons';
+
+const { matchFlavours } = BlocksUtils;
+
+type Selections = {
+  text?: TextSelection;
+  blocks?: BlockSelection[];
+  images?: ImageSelection[];
+};
+
+export type ChatAction = {
+  icon: TemplateResult<1>;
+  title: string;
+  toast: string;
+  showWhen: (host: EditorHost) => boolean;
+  handler: (
+    host: EditorHost,
+    content: string,
+    currentSelections: Selections,
+    chatSessionId?: string,
+    messageId?: string
+  ) => Promise<boolean>;
+};
+
+export async function queryHistoryMessages(doc: Doc, forkSessionId: string) {
+  // Get fork session messages
+  const histories = await AIProvider.histories?.chats(
+    doc.collection.id,
+    doc.id,
+    {
+      sessionId: forkSessionId,
+      messageOrder: ChatHistoryOrder.asc,
+    }
+  );
+
+  if (!histories || !histories.length) {
+    return [];
+  }
+
+  return histories[0].messages;
+}
+
+// Construct user info with messages
+export function constructUserInfoWithMessages(
+  messages: ChatMessage[],
+  userInfo: AIUserInfo | null
+) {
+  return messages.map(message => {
+    const { role, id, content, createdAt } = message;
+    const isUser = role === 'user';
+    const userInfoProps = isUser
+      ? {
+          userId: userInfo?.id,
+          userName: userInfo?.name,
+          avatarUrl: userInfo?.avatarUrl ?? undefined,
+        }
+      : {};
+    return {
+      id,
+      role,
+      content,
+      createdAt,
+      attachments: [],
+      ...userInfoProps,
+    };
+  });
+}
+
+export async function constructRootChatBlockMessages(
+  doc: Doc,
+  forkSessionId: string
+) {
+  // Convert chat messages to AI chat block messages
+  const userInfo = await AIProvider.userInfo;
+  const forkMessages = await queryHistoryMessages(doc, forkSessionId);
+  return constructUserInfoWithMessages(forkMessages, userInfo);
+}
+
+function getViewportCenter(
+  mode: DocMode,
+  rootService: PageRootService | EdgelessRootService
+) {
+  const center = { x: 0, y: 0 };
+  if (mode === 'page') {
+    const viewport = rootService.editPropsStore.getStorage('viewport');
+    if (viewport) {
+      if ('xywh' in viewport) {
+        const bound = Bound.deserialize(viewport.xywh);
+        center.x = bound.x + bound.w / 2;
+        center.y = bound.y + bound.h / 2;
+      } else {
+        center.x = viewport.centerX;
+        center.y = viewport.centerY;
+      }
+    }
+  } else {
+    // Else we should get latest viewport center from the edgeless root service
+    const edgelessService = rootService as EdgelessRootService;
+    center.x = edgelessService.viewport.centerX;
+    center.y = edgelessService.viewport.centerY;
+  }
+
+  return center;
+}
+
+// Add AI chat block and focus on it
+function addAIChatBlock(
+  doc: Doc,
+  messages: ChatMessage[],
+  sessionId: string,
+  viewportCenter: { x: number; y: number }
+) {
+  if (!messages.length || !sessionId) {
+    return;
+  }
+
+  const surfaceBlock = doc
+    .getBlocks()
+    .find(block => block.flavour === 'affine:surface');
+  if (!surfaceBlock) {
+    return;
+  }
+
+  // Add AI chat block to the center of the viewport
+  const width = 300; // AI_CHAT_BLOCK_WIDTH = 300
+  const height = 320; // AI_CHAT_BLOCK_HEIGHT = 320
+  const x = viewportCenter.x - width / 2;
+  const y = viewportCenter.y - height / 2;
+  const bound = new Bound(x, y, width, height);
+  const aiChatBlockId = doc.addBlock(
+    'affine:embed-ai-chat' as keyof BlockSuite.BlockModels,
+    {
+      xywh: bound.serialize(),
+      messages: JSON.stringify(messages),
+      sessionId,
+    },
+    surfaceBlock.id
+  );
+
+  return aiChatBlockId;
+}
+
+export function promptDocTitle(host: EditorHost, autofill?: string) {
+  const notification =
+    host.std.spec.getService('affine:page').notificationService;
+  if (!notification) return Promise.resolve(undefined);
+
+  return notification.prompt({
+    title: 'Create linked doc',
+    message: 'Enter a title for the new doc.',
+    placeholder: 'Untitled',
+    autofill,
+    confirmText: 'Confirm',
+    cancelText: 'Cancel',
+  });
+}
+
+const REPLACE_SELECTION = {
+  icon: ReplaceIcon,
+  title: 'Replace selection',
+  showWhen: (host: EditorHost) => {
+    const textSelection = host.selection.find('text');
+    const blockSelections = host.selection.filter('block');
+    if (
+      (!textSelection || textSelection.from.length === 0) &&
+      blockSelections?.length === 0
+    ) {
+      return false;
+    }
+    return true;
+  },
+  toast: 'Successfully replaced',
+  handler: async (
+    host: EditorHost,
+    content: string,
+    currentSelections: Selections
+  ) => {
+    const currentTextSelection = currentSelections.text;
+    const currentBlockSelections = currentSelections.blocks;
+    const [_, data] = host.command
+      .chain()
+      .getSelectedBlocks({
+        currentTextSelection,
+        currentBlockSelections,
+      })
+      .run();
+    if (!data.selectedBlocks) return false;
+
+    reportResponse('result:replace');
+
+    if (currentTextSelection) {
+      const { doc } = host;
+      const block = doc.getBlock(currentTextSelection.blockId);
+      if (matchFlavours(block?.model ?? null, ['affine:paragraph'])) {
+        block?.model.text?.replace(
+          currentTextSelection.from.index,
+          currentTextSelection.from.length,
+          content
+        );
+        return true;
+      }
+    }
+
+    await replace(
+      host,
+      content,
+      data.selectedBlocks[0],
+      data.selectedBlocks.map(block => block.model),
+      currentTextSelection
+    );
+    return true;
+  },
+};
+
+const INSERT_BELOW = {
+  icon: InsertBelowIcon,
+  title: 'Insert below',
+  showWhen: () => true,
+  toast: 'Successfully inserted',
+  handler: async (
+    host: EditorHost,
+    content: string,
+    currentSelections: Selections
+  ) => {
+    const currentTextSelection = currentSelections.text;
+    const currentBlockSelections = currentSelections.blocks;
+    const currentImageSelections = currentSelections.images;
+    const [_, data] = host.command
+      .chain()
+      .getSelectedBlocks({
+        currentTextSelection,
+        currentBlockSelections,
+        currentImageSelections,
+      })
+      .run();
+    if (!data.selectedBlocks) return false;
+    reportResponse('result:insert');
+    await insertBelow(
+      host,
+      content,
+      data.selectedBlocks[data.selectedBlocks?.length - 1]
+    );
+    return true;
+  },
+};
+
+const SAVE_CHAT_TO_BLOCK_ACTION: ChatAction = {
+  icon: BlockIcon,
+  title: 'Save chat to block',
+  toast: 'Successfully saved chat to a block',
+  showWhen: (host: EditorHost) =>
+    !!host.doc.awarenessStore.getFlag('enable_ai_chat_block'),
+  handler: async (
+    host: EditorHost,
+    _,
+    __,
+    chatSessionId?: string,
+    messageId?: string
+  ) => {
+    // The chat session id and the latest message id are required to fork the chat session
+    const parentSessionId = chatSessionId;
+    if (!messageId || !parentSessionId) {
+      return false;
+    }
+
+    const rootService = host.spec.getService('affine:page');
+    if (!rootService) return false;
+
+    const { docModeService, notificationService } = rootService;
+    const curMode = docModeService.getMode();
+    const viewportCenter = getViewportCenter(curMode, rootService);
+    // If current mode is not edgeless, switch to edgeless mode first
+    if (curMode !== 'edgeless') {
+      // Set mode to edgeless
+      docModeService.setMode('edgeless');
+      // Notify user to switch to edgeless mode
+      notificationService?.notify({
+        title: 'Save chat to a block',
+        accent: 'info',
+        message:
+          'This feature is not available in the page editor. Switch to edgeless mode.',
+        onClose: function (): void {},
+      });
+    }
+
+    try {
+      const newSessionId = await AIProvider.forkChat?.({
+        workspaceId: host.doc.collection.id,
+        docId: host.doc.id,
+        sessionId: parentSessionId,
+        latestMessageId: messageId,
+      });
+
+      if (!newSessionId) {
+        return false;
+      }
+
+      // Get messages before the latest message
+      const messages = await constructRootChatBlockMessages(
+        host.doc,
+        newSessionId
+      );
+
+      // After switching to edgeless mode, the user can save the chat to a block
+      const blockId = addAIChatBlock(
+        host.doc,
+        messages,
+        newSessionId,
+        viewportCenter
+      );
+      if (!blockId) {
+        return false;
+      }
+      return true;
+    } catch (err) {
+      console.error(err);
+      notificationService?.notify({
+        title: 'Failed to save chat to a block',
+        accent: 'error',
+        onClose: function (): void {},
+      });
+      return false;
+    }
+  },
+};
+
+const ADD_TO_EDGELESS_AS_NOTE = {
+  icon: CreateIcon,
+  title: 'Add to edgeless as note',
+  showWhen: () => true,
+  toast: 'New note created',
+  handler: async (host: EditorHost, content: string) => {
+    reportResponse('result:add-note');
+    const { doc } = host;
+    const service = host.spec.getService<EdgelessRootService>('affine:page');
+    const elements = service.selection.selectedElements;
+
+    const props: { displayMode: NoteDisplayMode; xywh?: SerializedXYWH } = {
+      displayMode: NoteDisplayMode.EdgelessOnly,
+    };
+
+    if (elements.length > 0) {
+      const bound = getElementsBound(
+        elements.map(e => Bound.deserialize(e.xywh))
+      );
+      const newBound = new Bound(bound.x, bound.maxY + 10, bound.w);
+      props.xywh = newBound.serialize();
+    }
+
+    const id = doc.addBlock('affine:note', props, doc.root?.id);
+
+    await insertFromMarkdown(host, content, doc, id, 0);
+
+    service.selection.set({
+      elements: [id],
+      editing: false,
+    });
+
+    return true;
+  },
+};
+
+const CREATE_AS_DOC = {
+  icon: CreateIcon,
+  title: 'Create as a doc',
+  showWhen: () => true,
+  toast: 'New doc created',
+  handler: (host: EditorHost, content: string) => {
+    reportResponse('result:add-page');
+    const newDoc = host.doc.collection.createDoc();
+    newDoc.load();
+    const rootId = newDoc.addBlock('affine:page');
+    newDoc.addBlock('affine:surface', {}, rootId);
+    const noteId = newDoc.addBlock('affine:note', {}, rootId);
+
+    host.spec.getService('affine:page').slots.docLinkClicked.emit({
+      docId: newDoc.id,
+    });
+    let complete = false;
+    (function addContent() {
+      if (complete) return;
+      const newHost = document.querySelector('editor-host');
+      // FIXME: this is a hack to wait for the host to be ready, now we don't have a way to know if the new host is ready
+      if (!newHost || newHost === host) {
+        setTimeout(addContent, 100);
+        return;
+      }
+      complete = true;
+      const { doc } = newHost;
+      insertFromMarkdown(newHost, content, doc, noteId, 0).catch(console.error);
+    })();
+
+    return true;
+  },
+};
+
+const CREATE_AS_LINKED_DOC = {
+  icon: CreateIcon,
+  title: 'Create as a linked doc',
+  showWhen: () => true,
+  toast: 'New doc created',
+  handler: async (host: EditorHost, content: string) => {
+    reportResponse('result:add-page');
+
+    const { doc } = host;
+    const surfaceBlock = doc
+      .getBlocks()
+      .find(block => block.flavour === 'affine:surface');
+    if (!surfaceBlock) {
+      return false;
+    }
+
+    const service = host.spec.getService<EdgelessRootService>('affine:page');
+    const mode = service.docModeService.getMode();
+    if (mode !== 'edgeless') {
+      return false;
+    }
+
+    // Create a new doc and add the content to it
+    const newDoc = host.doc.collection.createDoc();
+    newDoc.load();
+    const rootId = newDoc.addBlock('affine:page');
+    newDoc.addBlock('affine:surface', {}, rootId);
+    const noteId = newDoc.addBlock('affine:note', {}, rootId);
+    await insertFromMarkdown(host, content, newDoc, noteId, 0);
+
+    // Add a linked doc card to link to the new doc
+    const elements = service.selection.selectedElements;
+    const width = 364;
+    const height = 390;
+    let x = 0;
+    let y = 0;
+    if (elements.length) {
+      // Calculate the bound of the selected elements first
+      const bound = getElementsBound(
+        elements.map(e => Bound.deserialize(e.xywh))
+      );
+      x = bound.x;
+      y = bound.y + bound.h + 100;
+    }
+
+    // If the selected elements are not in the viewport, center the linked doc card
+    if (x === Number.POSITIVE_INFINITY || y === Number.POSITIVE_INFINITY) {
+      const viewportCenter = getViewportCenter(mode, service);
+      x = viewportCenter.x - width / 2;
+      y = viewportCenter.y - height / 2;
+    }
+
+    service.addBlock(
+      'affine:embed-linked-doc',
+      {
+        xywh: `[${x}, ${y}, ${width}, ${height}]`,
+        style: 'vertical',
+        pageId: newDoc.id,
+      },
+      surfaceBlock.id
+    );
+
+    return true;
+  },
+};
+
+const CommonActions: ChatAction[] = [REPLACE_SELECTION, INSERT_BELOW];
+
+export const PageEditorActions = [
+  ...CommonActions,
+  CREATE_AS_DOC,
+  SAVE_CHAT_TO_BLOCK_ACTION,
+];
+
+export const EdgelessEditorActions = [
+  ...CommonActions,
+  ADD_TO_EDGELESS_AS_NOTE,
+  SAVE_CHAT_TO_BLOCK_ACTION,
+];
+
+export const ChatBlockPeekViewActions = [
+  ADD_TO_EDGELESS_AS_NOTE,
+  CREATE_AS_LINKED_DOC,
+];

packages/frontend/core/src/blocksuite/presets/ai/_common/components/chat-action-list.ts

@@ -0,0 +1,171 @@
+import type {
+  BlockSelection,
+  EditorHost,
+  TextSelection,
+} from '@blocksuite/block-std';
+import type { ImageSelection } from '@blocksuite/blocks';
+import { css, html, LitElement, nothing } from 'lit';
+import { customElement, property } from 'lit/decorators.js';
+import { classMap } from 'lit/directives/class-map.js';
+import { repeat } from 'lit/directives/repeat.js';
+
+import { insertBelow } from '../../utils/editor-actions';
+import type { ChatAction } from '../chat-actions-handle';
+
+@customElement('chat-action-list')
+export class ChatActionList extends LitElement {
+  static override styles = css`
+    .actions-container {
+      display: flex;
+      gap: 8px;
+    }
+    .actions-container > div {
+      display: flex;
+      gap: 8px;
+    }
+    .actions-container.horizontal {
+      flex-wrap: wrap;
+      justify-content: end;
+    }
+    .actions-container.vertical {
+      flex-direction: column;
+      align-items: flex-end;
+    }
+    .action {
+      width: fit-content;
+      height: 32px;
+      padding: 12px;
+      box-sizing: border-box;
+      border-radius: 8px;
+      border: 1px solid var(--affine-border-color);
+      background-color: var(--affine-white-10);
+      display: flex;
+      flex-direction: row;
+      align-items: center;
+      gap: 4px;
+      font-size: 15px;
+      font-weight: 500;
+      color: var(--affine-text-primary-color);
+      cursor: pointer;
+      user-select: none;
+    }
+    .action svg {
+      color: var(--affine-icon-color);
+    }
+  `;
+
+  private get _selectionValue() {
+    return this.host.selection.value;
+  }
+
+  private get _rootService() {
+    return this.host.spec.getService('affine:page');
+  }
+
+  private get _currentTextSelection(): TextSelection | undefined {
+    return this._selectionValue.find(v => v.type === 'text') as TextSelection;
+  }
+
+  private get _currentBlockSelections(): BlockSelection[] | undefined {
+    return this._selectionValue.filter(v => v.type === 'block');
+  }
+
+  private get _currentImageSelections(): ImageSelection[] | undefined {
+    return this._selectionValue.filter(v => v.type === 'image');
+  }
+
+  @property({ attribute: false })
+  accessor host!: EditorHost;
+
+  @property({ attribute: false })
+  accessor actions: ChatAction[] = [];
+
+  @property({ attribute: false })
+  accessor content: string = '';
+
+  @property({ attribute: false })
+  accessor chatSessionId: string | undefined = undefined;
+
+  @property({ attribute: false })
+  accessor messageId: string | undefined = undefined;
+
+  @property({ attribute: false })
+  accessor layoutDirection: 'horizontal' | 'vertical' = 'vertical'; // New property for layout direction
+
+  @property({ attribute: false })
+  accessor withMargin = false;
+
+  override render() {
+    const { actions } = this;
+    if (!actions.length) {
+      return nothing;
+    }
+
+    const { host, content, chatSessionId, messageId, layoutDirection } = this;
+    const classes = classMap({
+      'actions-container': true,
+      horizontal: layoutDirection === 'horizontal',
+      vertical: layoutDirection === 'vertical',
+    });
+
+    return html`<style>
+        .actions-container {
+          margin-top: ${this.withMargin ? '8px' : '0'};
+        }
+      </style>
+      <div class=${classes}>
+        ${repeat(
+          actions.filter(action => action.showWhen(host)),
+          action => action.title,
+          action => {
+            return html`<div class="action">
+              ${action.icon}
+              <div
+                @click=${async () => {
+                  if (
+                    action.title === 'Insert below' &&
+                    this._selectionValue.length === 1 &&
+                    this._selectionValue[0].type === 'database'
+                  ) {
+                    const element = this.host.view.getBlock(
+                      this._selectionValue[0].blockId
+                    );
+                    if (!element) return;
+                    await insertBelow(host, content, element);
+                    return;
+                  }
+                  const currentSelections = {
+                    text: this._currentTextSelection,
+                    blocks: this._currentBlockSelections,
+                    images: this._currentImageSelections,
+                  };
+                  const success = await action.handler(
+                    host,
+                    content,
+                    currentSelections,
+                    chatSessionId,
+                    messageId
+                  );
+                  if (success) {
+                    this._rootService.notificationService?.notify({
+                      title: action.toast,
+                      accent: 'success',
+                      onClose: function (): void {},
+                    });
+                  }
+                }}
+              >
+                ${action.title}
+              </div>
+            </div>`;
+          }
+        )}
+      </div>`;
+  }
+}
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'chat-action-list': ChatActionList;
+  }
+}

packages/frontend/core/src/blocksuite/presets/ai/_common/components/copy-more.ts

@@ -4,17 +4,15 @@ import type {
   TextSelection,
 } from '@blocksuite/block-std';
 import { WithDisposable } from '@blocksuite/block-std';
-import { type AIError, createButtonPopper, Tooltip } from '@blocksuite/blocks';
+import { createButtonPopper, Tooltip } from '@blocksuite/blocks';
 import { noop } from '@blocksuite/global/utils';
 import { css, html, LitElement, nothing, type PropertyValues } from 'lit';
 import { customElement, property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
 
+import { type ChatAction } from '../../_common/chat-actions-handle';
 import { CopyIcon, MoreIcon, RetryIcon } from '../../_common/icons';
-import { AIProvider } from '../../provider';
 import { copyText } from '../../utils/editor-actions';
-import type { ChatContextValue, ChatMessage } from '../chat-context';
-import { PageEditorActions } from './actions-handle';
 
 noop(Tooltip);
 
@@ -25,10 +23,10 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
       display: flex;
       gap: 8px;
       height: 36px;
+      box-sizing: border-box;
       justify-content: flex-end;
       align-items: center;
-      margin-top: 8px;
-      margin-bottom: 12px;
+      padding: 8px 0;
 
       div {
         cursor: pointer;
@@ -74,6 +72,22 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
     }
   `;
 
+  private get _rootService() {
+    return this.host.spec.getService('affine:page');
+  }
+
+  private get _selectionValue() {
+    return this.host.selection.value;
+  }
+
+  private get _currentTextSelection(): TextSelection | undefined {
+    return this._selectionValue.find(v => v.type === 'text') as TextSelection;
+  }
+
+  private get _currentBlockSelections(): BlockSelection[] | undefined {
+    return this._selectionValue.filter(v => v.type === 'block');
+  }
+
   @state()
   private accessor _showMoreMenu = false;
 
@@ -88,34 +102,33 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
   @property({ attribute: false })
   accessor host!: EditorHost;
 
+  @property({ attribute: false })
+  accessor actions: ChatAction[] = [];
+
   @property({ attribute: false })
   accessor content!: string;
 
   @property({ attribute: false })
-  accessor messageId!: string;
+  accessor chatSessionId: string | undefined = undefined;
+
+  @property({ attribute: false })
+  accessor messageId: string | undefined = undefined;
 
   @property({ attribute: false })
   accessor isLast!: boolean;
 
   @property({ attribute: false })
-  accessor curTextSelection: TextSelection | undefined = undefined;
+  accessor withMargin = false;
 
   @property({ attribute: false })
-  accessor curBlockSelections: BlockSelection[] | undefined = undefined;
-
-  @property({ attribute: false })
-  accessor chatContextValue!: ChatContextValue;
-
-  @property({ attribute: false })
-  accessor updateContext!: (context: Partial<ChatContextValue>) => void;
+  accessor retry = () => {};
 
   private _toggle() {
     this._morePopper?.toggle();
   }
 
   private readonly _notifySuccess = (title: string) => {
-    const rootService = this.host.spec.getService('affine:page');
-    const { notificationService } = rootService;
+    const { notificationService } = this._rootService;
     notificationService?.notify({
       title: title,
       accent: 'success',
@@ -123,48 +136,6 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
     });
   };
 
-  private async _retry() {
-    const { doc } = this.host;
-    try {
-      const abortController = new AbortController();
-
-      const items = [...this.chatContextValue.items];
-      const last = items[items.length - 1];
-      if ('content' in last) {
-        last.content = '';
-        last.createdAt = new Date().toISOString();
-      }
-      this.updateContext({ items, status: 'loading', error: null });
-
-      const stream = AIProvider.actions.chat?.({
-        retry: true,
-        docId: doc.id,
-        workspaceId: doc.collection.id,
-        host: this.host,
-        stream: true,
-        signal: abortController.signal,
-        where: 'chat-panel',
-        control: 'chat-send',
-      });
-
-      if (stream) {
-        this.updateContext({ abortController });
-        for await (const text of stream) {
-          const items = [...this.chatContextValue.items];
-          const last = items[items.length - 1] as ChatMessage;
-          last.content += text;
-          this.updateContext({ items, status: 'transmitting' });
-        }
-
-        this.updateContext({ status: 'success' });
-      }
-    } catch (error) {
-      this.updateContext({ status: 'error', error: error as AIError });
-    } finally {
-      this.updateContext({ abortController: null });
-    }
-  }
-
   protected override updated(changed: PropertyValues): void {
     if (changed.has('isLast')) {
       if (this.isLast) {
@@ -185,8 +156,12 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
   }
 
   override render() {
-    const { host, content, isLast, messageId } = this;
+    const { host, content, isLast, messageId, chatSessionId, actions } = this;
     return html`<style>
+        .copy-more {
+          margin-top: ${this.withMargin ? '8px' : '0px'};
+          margin-bottom: ${this.withMargin ? '12px' : '0px'};
+        }
         .more-menu {
           padding: ${this._showMoreMenu ? '8px' : '0px'};
         }
@@ -206,7 +181,7 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
             </div>`
           : nothing}
         ${isLast
-          ? html`<div @click=${() => this._retry()}>
+          ? html`<div @click=${() => this.retry()}>
               ${RetryIcon}
               <affine-tooltip>Retry</affine-tooltip>
             </div>`
@@ -221,12 +196,12 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
       <div class="more-menu">
         ${this._showMoreMenu
           ? repeat(
-              PageEditorActions.filter(action => action.showWhen(host)),
+              actions.filter(action => action.showWhen(host)),
               action => action.title,
               action => {
                 const currentSelections = {
-                  text: this.curTextSelection,
-                  blocks: this.curBlockSelections,
+                  text: this._currentTextSelection,
+                  blocks: this._currentBlockSelections,
                 };
                 return html`<div
                   @click=${async () => {
@@ -234,8 +209,8 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
                       host,
                       content,
                       currentSelections,
-                      this.chatContextValue,
-                      messageId ?? undefined
+                      chatSessionId,
+                      messageId
                     );
 
                     if (success) {

packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-handler.ts

@@ -15,6 +15,7 @@ import {
   TextElementModel,
 } from '@blocksuite/blocks';
 import { assertExists } from '@blocksuite/global/utils';
+import { AIChatBlockModel } from '@blocksuite/presets';
 import { Slice } from '@blocksuite/store';
 import type { TemplateResult } from 'lit';
 
@@ -553,3 +554,18 @@ export function mindmapRootShowWhen(_: unknown, __: unknown, host: EditorHost) {
 
   return selected.length === 1 && isMindMapRoot(selected[0]);
 }
+
+// TODO(@chen): remove this function after the new AI chat block related function is fully implemented
+export function notAllAIChatBlockShowWhen(
+  _: unknown,
+  __: unknown,
+  host: EditorHost
+) {
+  const selected = getCopilotSelectedElems(host);
+  if (selected.length === 0) return true;
+
+  const allAIChatBlocks = selected.every(
+    model => model instanceof AIChatBlockModel
+  );
+  return !allAIChatBlocks;
+}

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

@@ -50,7 +50,11 @@ declare global {
       | 'chat-send'
       | 'block-action-bar';
 
-    type TrackerWhere = 'chat-panel' | 'inline-chat-panel' | 'ai-panel';
+    type TrackerWhere =
+      | 'chat-panel'
+      | 'inline-chat-panel'
+      | 'ai-panel'
+      | 'ai-chat-block';
 
     interface TrackerOptions {
       control: TrackerControl;
@@ -102,6 +106,11 @@ declare global {
     type AIActionTextResponse<T extends AITextActionOptions> =
       T['stream'] extends true ? TextStream : Promise<string>;
 
+    interface ChatOptions extends AITextActionOptions {
+      sessionId?: string;
+      isRootSession?: boolean;
+    }
+
     interface TranslateOptions extends AITextActionOptions {
       lang: (typeof translateLangs)[number];
     }
@@ -120,7 +129,7 @@ declare global {
 
     interface AIActions {
       // chat is a bit special because it's has a internally maintained session
-      chat<T extends AITextActionOptions>(options: T): AIActionTextResponse<T>;
+      chat<T extends ChatOptions>(options: T): AIActionTextResponse<T>;
 
       summary<T extends AITextActionOptions>(
         options: T
@@ -234,6 +243,7 @@ declare global {
         content: string;
         createdAt: string;
         role: MessageRole;
+        attachments?: string[];
       }[];
     }
 

packages/frontend/core/src/blocksuite/presets/ai/ai-panel.ts

@@ -119,7 +119,7 @@ function createNewNote(host: EditorHost): AIItemConfig {
         );
 
         assertExists(panel.answer);
-        insertFromMarkdown(host, panel.answer, noteBlockId)
+        insertFromMarkdown(host, panel.answer, doc, noteBlockId)
           .then(() => {
             service.selection.set({
               elements: [noteBlockId],

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

@@ -1,392 +0,0 @@
-import { ChatHistoryOrder } from '@affine/graphql';
-import type {
-  BlockSelection,
-  EditorHost,
-  TextSelection,
-} from '@blocksuite/block-std';
-import type {
-  DocMode,
-  EdgelessRootService,
-  ImageSelection,
-  PageRootService,
-} from '@blocksuite/blocks';
-import {
-  BlocksUtils,
-  getElementsBound,
-  NoteDisplayMode,
-} from '@blocksuite/blocks';
-import type { SerializedXYWH } from '@blocksuite/global/utils';
-import { Bound } from '@blocksuite/global/utils';
-import type { Doc } from '@blocksuite/store';
-import type { TemplateResult } from 'lit';
-
-import {
-  BlockIcon,
-  CreateIcon,
-  InsertBelowIcon,
-  ReplaceIcon,
-} from '../../_common/icons';
-import { AIProvider } from '../../provider';
-import { reportResponse } from '../../utils/action-reporter';
-import { insertBelow, replace } from '../../utils/editor-actions';
-import { insertFromMarkdown } from '../../utils/markdown-utils';
-import type { ChatBlockMessage, ChatContextValue } from '../chat-context';
-
-const { matchFlavours } = BlocksUtils;
-
-type Selections = {
-  text?: TextSelection;
-  blocks?: BlockSelection[];
-  images?: ImageSelection[];
-};
-
-type ChatAction = {
-  icon: TemplateResult<1>;
-  title: string;
-  toast: string;
-  showWhen: (host: EditorHost) => boolean;
-  handler: (
-    host: EditorHost,
-    content: string,
-    currentSelections: Selections,
-    chatContext?: ChatContextValue,
-    messageId?: string
-  ) => Promise<boolean>;
-};
-
-async function constructChatBlockMessages(doc: Doc, forkSessionId: string) {
-  const userInfo = await AIProvider.userInfo;
-  // Get fork session messages
-  const histories = await AIProvider.histories?.chats(
-    doc.collection.id,
-    doc.id,
-    {
-      sessionId: forkSessionId,
-      messageOrder: ChatHistoryOrder.asc,
-    }
-  );
-
-  if (!histories || !histories.length) {
-    return [];
-  }
-
-  const messages = histories[0].messages.map(message => {
-    const { role, id, content, createdAt } = message;
-    const isUser = role === 'user';
-    const userInfoProps = isUser
-      ? {
-          userId: userInfo?.id,
-          userName: userInfo?.name,
-          avatarUrl: userInfo?.avatarUrl ?? undefined,
-        }
-      : {};
-    return {
-      id,
-      role,
-      content,
-      createdAt,
-      attachments: [],
-      ...userInfoProps,
-    };
-  });
-  return messages;
-}
-
-function getViewportCenter(
-  mode: DocMode,
-  rootService: PageRootService | EdgelessRootService
-) {
-  const center = { x: 0, y: 0 };
-  if (mode === 'page') {
-    const viewport = rootService.editPropsStore.getStorage('viewport');
-    if (viewport) {
-      if ('xywh' in viewport) {
-        const bound = Bound.deserialize(viewport.xywh);
-        center.x = bound.x + bound.w / 2;
-        center.y = bound.y + bound.h / 2;
-      } else {
-        center.x = viewport.centerX;
-        center.y = viewport.centerY;
-      }
-    }
-  } else {
-    // Else we should get latest viewport center from the edgeless root service
-    const edgelessService = rootService as EdgelessRootService;
-    center.x = edgelessService.viewport.centerX;
-    center.y = edgelessService.viewport.centerY;
-  }
-
-  return center;
-}
-
-// Add AI chat block and focus on it
-function addAIChatBlock(
-  doc: Doc,
-  messages: ChatBlockMessage[],
-  sessionId: string,
-  viewportCenter: { x: number; y: number }
-) {
-  if (!messages.length || !sessionId) {
-    return;
-  }
-
-  const surfaceBlock = doc
-    .getBlocks()
-    .find(block => block.flavour === 'affine:surface');
-  if (!surfaceBlock) {
-    return;
-  }
-
-  // Add AI chat block to the center of the viewport
-  const width = 300; // AI_CHAT_BLOCK_WIDTH = 300
-  const height = 320; // AI_CHAT_BLOCK_HEIGHT = 320
-  const x = viewportCenter.x - width / 2;
-  const y = viewportCenter.y - height / 2;
-  const bound = new Bound(x, y, width, height);
-  const aiChatBlockId = doc.addBlock(
-    'affine:embed-ai-chat' as keyof BlockSuite.BlockModels,
-    {
-      xywh: bound.serialize(),
-      messages: JSON.stringify(messages),
-      sessionId,
-    },
-    surfaceBlock.id
-  );
-
-  return aiChatBlockId;
-}
-
-const CommonActions: ChatAction[] = [
-  {
-    icon: ReplaceIcon,
-    title: 'Replace selection',
-    showWhen: () => true,
-    toast: 'Successfully replaced',
-    handler: async (
-      host: EditorHost,
-      content: string,
-      currentSelections: Selections
-    ) => {
-      const currentTextSelection = currentSelections.text;
-      const currentBlockSelections = currentSelections.blocks;
-      const [_, data] = host.command
-        .chain()
-        .getSelectedBlocks({
-          currentTextSelection,
-          currentBlockSelections,
-        })
-        .run();
-      if (!data.selectedBlocks) return false;
-
-      reportResponse('result:replace');
-
-      if (currentTextSelection) {
-        const { doc } = host;
-        const block = doc.getBlock(currentTextSelection.blockId);
-        if (matchFlavours(block?.model ?? null, ['affine:paragraph'])) {
-          block?.model.text?.replace(
-            currentTextSelection.from.index,
-            currentTextSelection.from.length,
-            content
-          );
-          return true;
-        }
-      }
-
-      await replace(
-        host,
-        content,
-        data.selectedBlocks[0],
-        data.selectedBlocks.map(block => block.model),
-        currentTextSelection
-      );
-      return true;
-    },
-  },
-  {
-    icon: InsertBelowIcon,
-    title: 'Insert below',
-    showWhen: () => true,
-    toast: 'Successfully inserted',
-    handler: async (
-      host: EditorHost,
-      content: string,
-      currentSelections: Selections
-    ) => {
-      const currentTextSelection = currentSelections.text;
-      const currentBlockSelections = currentSelections.blocks;
-      const currentImageSelections = currentSelections.images;
-      const [_, data] = host.command
-        .chain()
-        .getSelectedBlocks({
-          currentTextSelection,
-          currentBlockSelections,
-          currentImageSelections,
-        })
-        .run();
-      if (!data.selectedBlocks) return false;
-      reportResponse('result:insert');
-      await insertBelow(
-        host,
-        content,
-        data.selectedBlocks[data.selectedBlocks?.length - 1]
-      );
-      return true;
-    },
-  },
-];
-
-const SAVE_CHAT_TO_BLOCK_ACTION: ChatAction = {
-  icon: BlockIcon,
-  title: 'Save chat to block',
-  toast: 'Successfully saved chat to a block',
-  showWhen: (host: EditorHost) =>
-    !!host.doc.awarenessStore.getFlag('enable_ai_chat_block'),
-  handler: async (
-    host: EditorHost,
-    _,
-    __,
-    chatContext?: ChatContextValue,
-    messageId?: string
-  ) => {
-    // The chat session id and the latest message id are required to fork the chat session
-    const parentSessionId = chatContext?.chatSessionId;
-    if (!messageId || !parentSessionId) {
-      return false;
-    }
-
-    const rootService = host.spec.getService('affine:page');
-    if (!rootService) return false;
-
-    const { docModeService, notificationService } = rootService;
-    const curMode = docModeService.getMode();
-    const viewportCenter = getViewportCenter(curMode, rootService);
-    // If current mode is not edgeless, switch to edgeless mode first
-    if (curMode !== 'edgeless') {
-      // Set mode to edgeless
-      docModeService.setMode('edgeless');
-      // Notify user to switch to edgeless mode
-      notificationService?.notify({
-        title: 'Save chat to a block',
-        accent: 'info',
-        message:
-          'This feature is not available in the page editor. Switch to edgeless mode.',
-        onClose: function (): void {},
-      });
-    }
-
-    try {
-      const newSessionId = await AIProvider.forkChat?.({
-        workspaceId: host.doc.collection.id,
-        docId: host.doc.id,
-        sessionId: parentSessionId,
-        latestMessageId: messageId,
-      });
-
-      if (!newSessionId) {
-        return false;
-      }
-
-      // Construct chat block messages from the forked chat session
-      const messages = await constructChatBlockMessages(host.doc, newSessionId);
-
-      // After switching to edgeless mode, the user can save the chat to a block
-      const blockId = addAIChatBlock(
-        host.doc,
-        messages,
-        newSessionId,
-        viewportCenter
-      );
-      if (!blockId) {
-        return false;
-      }
-
-      return true;
-    } catch (err) {
-      console.error(err);
-      notificationService?.notify({
-        title: 'Failed to save chat to a block',
-        accent: 'error',
-        onClose: function (): void {},
-      });
-      return false;
-    }
-  },
-};
-
-export const PageEditorActions = [
-  ...CommonActions,
-  {
-    icon: CreateIcon,
-    title: 'Create as a doc',
-    showWhen: () => true,
-    toast: 'New doc created',
-    handler: (host: EditorHost, content: string) => {
-      reportResponse('result:add-page');
-      const newDoc = host.doc.collection.createDoc();
-      newDoc.load();
-      const rootId = newDoc.addBlock('affine:page');
-      newDoc.addBlock('affine:surface', {}, rootId);
-      const noteId = newDoc.addBlock('affine:note', {}, rootId);
-
-      host.spec.getService('affine:page').slots.docLinkClicked.emit({
-        docId: newDoc.id,
-      });
-      let complete = false;
-      (function addContent() {
-        if (complete) return;
-        const newHost = document.querySelector('editor-host');
-        // FIXME: this is a hack to wait for the host to be ready, now we don't have a way to know if the new host is ready
-        if (!newHost || newHost === host) {
-          setTimeout(addContent, 100);
-          return;
-        }
-        complete = true;
-        insertFromMarkdown(newHost, content, noteId, 0).catch(console.error);
-      })();
-
-      return true;
-    },
-  },
-  SAVE_CHAT_TO_BLOCK_ACTION,
-];
-
-export const EdgelessEditorActions = [
-  ...CommonActions,
-  {
-    icon: CreateIcon,
-    title: 'Add to edgeless as note',
-    showWhen: () => true,
-    toast: 'New note created',
-    handler: async (host: EditorHost, content: string) => {
-      reportResponse('result:add-note');
-      const { doc } = host;
-      const service = host.spec.getService<EdgelessRootService>('affine:page');
-      const elements = service.selection.selectedElements;
-
-      const props: { displayMode: NoteDisplayMode; xywh?: SerializedXYWH } = {
-        displayMode: NoteDisplayMode.EdgelessOnly,
-      };
-
-      if (elements.length > 0) {
-        const bound = getElementsBound(
-          elements.map(e => Bound.deserialize(e.xywh))
-        );
-        const newBound = new Bound(bound.x, bound.maxY + 10, bound.w);
-        props.xywh = newBound.serialize();
-      }
-
-      const id = doc.addBlock('affine:note', props, doc.root?.id);
-
-      await insertFromMarkdown(host, content, id, 0);
-
-      service.selection.set({
-        elements: [id],
-        editing: false,
-      });
-
-      return true;
-    },
-  },
-  SAVE_CHAT_TO_BLOCK_ACTION,
-];

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

@@ -428,7 +428,6 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
 
     const content = (markdown ? `${markdown}\n` : '') + text;
 
-    // TODO: Should update message id especially for the assistant message
     this.updateContext({
       items: [
         ...this.chatContextValue.items,
@@ -460,6 +459,7 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
         signal: abortController.signal,
         where: 'chat-panel',
         control: 'chat-send',
+        isRootSession: true,
       });
 
       if (stream) {
@@ -476,7 +476,7 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
 
         if (!this.chatContextValue.chatSessionId) {
           this.updateContext({
-            chatSessionId: AIProvider.LAST_ACTION_SESSIONID,
+            chatSessionId: AIProvider.LAST_ROOT_SESSION_ID,
           });
         }
 

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

@@ -7,20 +7,16 @@ import './actions/make-real';
 import './actions/slides';
 import './actions/mindmap';
 import './actions/chat-text';
-import './actions/copy-more';
 import './actions/image-to-text';
 import './actions/image';
 import './chat-cards';
+import '../_common/components/chat-action-list';
+import '../_common/components/copy-more';
 
-import type {
-  BaseSelection,
-  BlockSelection,
-  EditorHost,
-  TextSelection,
-} from '@blocksuite/block-std';
+import type { BaseSelection, EditorHost } from '@blocksuite/block-std';
 import { ShadowlessElement, WithDisposable } from '@blocksuite/block-std';
-import type { ImageSelection } from '@blocksuite/blocks';
 import {
+  type AIError,
   isInsidePageEditor,
   PaymentRequiredError,
   UnauthorizedError,
@@ -30,39 +26,19 @@ import { customElement, property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
 import { styleMap } from 'lit/directives/style-map.js';
 
-import { AffineAvatarIcon, AffineIcon, DownArrowIcon } from '../_common/icons';
-import {
-  GeneralErrorRenderer,
-  PaymentRequiredErrorRenderer,
-} from '../messages/error';
-import { AIProvider } from '../provider';
-import { insertBelow } from '../utils/editor-actions';
 import {
   EdgelessEditorActions,
   PageEditorActions,
-} from './actions/actions-handle';
+} from '../_common/chat-actions-handle';
+import { AffineAvatarIcon, AffineIcon, DownArrowIcon } from '../_common/icons';
+import { AIChatErrorRenderer } from '../messages/error';
+import { AIProvider } from '../provider';
 import type { ChatContextValue, ChatItem, ChatMessage } from './chat-context';
 import { HISTORY_IMAGE_ACTIONS } from './const';
 import { AIPreloadConfig } from './preload-config';
 
 @customElement('chat-panel-messages')
 export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
-  private get _currentTextSelection(): TextSelection | undefined {
-    return this._selectionValue.find(v => v.type === 'text') as TextSelection;
-  }
-
-  private get _currentBlockSelections(): BlockSelection[] | undefined {
-    return this._selectionValue.filter(v => v.type === 'block');
-  }
-
-  private get _currentImageSelections(): ImageSelection[] | undefined {
-    return this._selectionValue.filter(v => v.type === 'image');
-  }
-
-  private get _rootService() {
-    return this.host.spec.getService('affine:page');
-  }
-
   static override styles = css`
     chat-panel-messages {
       position: relative;
@@ -328,35 +304,9 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
     );
   }
 
-  renderError() {
-    const { error } = this.chatContextValue;
-
-    if (error instanceof PaymentRequiredError) {
-      return PaymentRequiredErrorRenderer(this.host);
-    } else if (error instanceof UnauthorizedError) {
-      return GeneralErrorRenderer(
-        html`You need to login to AFFiNE Cloud to continue using AFFiNE AI.`,
-        html`<div
-          style=${styleMap({
-            padding: '4px 12px',
-            borderRadius: '8px',
-            border: '1px solid var(--affine-border-color)',
-            cursor: 'pointer',
-            backgroundColor: 'var(--affine-hover-color)',
-          })}
-          @click=${() =>
-            AIProvider.slots.requestLogin.emit({ host: this.host })}
-        >
-          Login
-        </div>`
-      );
-    } else {
-      return GeneralErrorRenderer();
-    }
-  }
-
   renderItem(item: ChatItem, isLast: boolean) {
     const { status, error } = this.chatContextValue;
+    const { host } = this;
 
     if (isLast && status === 'loading') {
       return this.renderLoading();
@@ -368,7 +318,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
       (error instanceof PaymentRequiredError ||
         error instanceof UnauthorizedError)
     ) {
-      return this.renderError();
+      return AIChatErrorRenderer(host, error);
     }
 
     if ('role' in item) {
@@ -377,48 +327,49 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
           ? 'finished'
           : 'generating'
         : 'finished';
+      const shouldRenderError = isLast && status === 'error' && !!error;
       return html`<chat-text
-          .host=${this.host}
+          .host=${host}
           .attachments=${item.attachments}
           .text=${item.content}
           .state=${state}
         ></chat-text>
-        ${isLast && status === 'error' ? this.renderError() : nothing}
+        ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
         ${this.renderEditorActions(item, isLast)}`;
     } else {
       switch (item.action) {
         case 'Create a presentation':
           return html`<action-slides
-            .host=${this.host}
+            .host=${host}
             .item=${item}
           ></action-slides>`;
         case 'Make it real':
           return html`<action-make-real
-            .host=${this.host}
+            .host=${host}
             .item=${item}
           ></action-make-real>`;
         case 'Brainstorm mindmap':
           return html`<action-mindmap
-            .host=${this.host}
+            .host=${host}
             .item=${item}
           ></action-mindmap>`;
         case 'Explain this image':
         case 'Generate a caption':
           return html`<action-image-to-text
-            .host=${this.host}
+            .host=${host}
             .item=${item}
           ></action-image-to-text>`;
         default:
           if (HISTORY_IMAGE_ACTIONS.includes(item.action)) {
             return html`<action-image
-              .host=${this.host}
+              .host=${host}
               .item=${item}
             ></action-image>`;
           }
 
           return html`<action-text
             .item=${item}
-            .host=${this.host}
+            .host=${host}
             .isCode=${item.action === 'Explain this code' ||
             item.action === 'Check code error'}
           ></action-text>`;
@@ -449,8 +400,54 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
     this.messagesContainer.scrollTo(0, this.messagesContainer.scrollHeight);
   }
 
+  retry = async () => {
+    const { doc } = this.host;
+    try {
+      const { chatSessionId } = this.chatContextValue;
+      if (!chatSessionId) return;
+
+      const abortController = new AbortController();
+      const items = [...this.chatContextValue.items];
+      const last = items[items.length - 1];
+      if ('content' in last) {
+        last.content = '';
+        last.createdAt = new Date().toISOString();
+      }
+      this.updateContext({ items, status: 'loading', error: null });
+
+      const stream = AIProvider.actions.chat?.({
+        sessionId: chatSessionId,
+        retry: true,
+        docId: doc.id,
+        workspaceId: doc.collection.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 items = [...this.chatContextValue.items];
+          const last = items[items.length - 1] as ChatMessage;
+          last.content += text;
+          this.updateContext({ items, status: 'transmitting' });
+        }
+
+        this.updateContext({ status: 'success' });
+      }
+    } catch (error) {
+      this.updateContext({ status: 'error', error: error as AIError });
+    } finally {
+      this.updateContext({ abortController: null });
+    }
+  };
+
   renderEditorActions(item: ChatMessage, isLast: boolean) {
-    const { status } = this.chatContextValue;
+    const { status, chatSessionId } = this.chatContextValue;
 
     if (item.role !== 'assistant') return nothing;
 
@@ -469,117 +466,25 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
       : EdgelessEditorActions;
 
     return html`
-      <style>
-        .actions-container {
-          display: flex;
-          flex-direction: column;
-          align-items: flex-end;
-          gap: 8px;
-          margin-top: 8px;
-        }
-
-        .actions-container > div {
-          display: flex;
-          gap: 8px;
-        }
-
-        .action {
-          width: fit-content;
-          height: 32px;
-          padding: 12px;
-          border-radius: 8px;
-          border: 1px solid var(--affine-border-color);
-          background-color: var(--affine-white-10);
-          display: flex;
-          flex-direction: row;
-          align-items: center;
-          gap: 4px;
-          font-size: 15px;
-          font-weight: 500;
-          color: var(--affine-text-primary-color);
-          cursor: pointer;
-          user-select: none;
-        }
-
-        .action svg {
-          color: var(--affine-icon-color);
-        }
-      </style>
       <chat-copy-more
         .host=${host}
+        .actions=${actions}
         .content=${content}
         .isLast=${isLast}
+        .chatSessionId=${chatSessionId ?? undefined}
         .messageId=${messageId}
-        .curTextSelection=${this._currentTextSelection}
-        .curBlockSelections=${this._currentBlockSelections}
-        .chatContextValue=${this.chatContextValue}
-        .updateContext=${this.updateContext}
+        .withMargin=${true}
+        .retry=${() => this.retry()}
       ></chat-copy-more>
-      ${isLast
-        ? html`<div class="actions-container">
-            ${repeat(
-              actions
-                .filter(action => action.showWhen(host))
-                .filter(action => {
-                  if (!content) return false;
-
-                  if (
-                    action.title === 'Replace selection' &&
-                    (!this._currentTextSelection ||
-                      this._currentTextSelection.from.length === 0) &&
-                    this._currentBlockSelections?.length === 0
-                  ) {
-                    return false;
-                  }
-                  return true;
-                }),
-              action => action.title,
-              action => {
-                return html`<div class="action">
-                  ${action.icon}
-                  <div
-                    @click=${async () => {
-                      if (
-                        action.title === 'Insert below' &&
-                        this._selectionValue.length === 1 &&
-                        this._selectionValue[0].type === 'database'
-                      ) {
-                        const element = this.host.view.getBlock(
-                          this._selectionValue[0].blockId
-                        );
-                        if (!element) return;
-                        await insertBelow(host, content, element);
-                        return;
-                      }
-
-                      const currentSelections = {
-                        text: this._currentTextSelection,
-                        blocks: this._currentBlockSelections,
-                        images: this._currentImageSelections,
-                      };
-
-                      const success = await action.handler(
-                        host,
-                        content,
-                        currentSelections,
-                        this.chatContextValue,
-                        messageId ?? undefined
-                      );
-                      if (success) {
-                        this._rootService.notificationService?.notify({
-                          title: action.toast,
-                          accent: 'success',
-                          onClose: function (): void {},
-                        });
-                      }
-                    }}
-                  >
-                    ${action.title}
-                  </div>
-                </div>`;
-              }
-            )}
-          </div>`
+      ${isLast && !!content
+        ? html`<chat-action-list
+            .actions=${actions}
+            .host=${host}
+            .content=${content}
+            .chatSessionId=${chatSessionId ?? undefined}
+            .messageId=${messageId ?? undefined}
+            .withMargin=${true}
+          ></chat-action-list>`
         : nothing}
     `;
   }

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

@@ -122,6 +122,7 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
         this._chatSessionId = history.sessionId;
         this.chatContextValue.chatSessionId = history.sessionId;
         items.push(...history.messages);
+        AIProvider.LAST_ROOT_SESSION_ID = history.sessionId;
       }
 
       this.chatContextValue = {
@@ -188,6 +189,7 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
 
   protected override updated(_changedProperties: PropertyValues) {
     if (_changedProperties.has('doc')) {
+      this.chatContextValue.chatSessionId = null;
       this._resetItems();
     }
   }

packages/frontend/core/src/blocksuite/presets/ai/entries/edgeless/actions-config.ts

@@ -35,6 +35,7 @@ import {
   imageOnlyShowWhen,
   mindmapChildShowWhen,
   mindmapRootShowWhen,
+  notAllAIChatBlockShowWhen,
   noteBlockOrTextShowWhen,
   noteWithCodeBlockShowWen,
 } from '../../actions/edgeless-handler';
@@ -271,7 +272,7 @@ const generateGroup: AIItemGroupConfig = {
     {
       name: 'Generate an image',
       icon: AIImageIcon,
-      showWhen: () => true,
+      showWhen: notAllAIChatBlockShowWhen,
       handler: actionToHandler(
         'createImage',
         AIImageIconWithAnimation,
@@ -403,7 +404,7 @@ const generateGroup: AIItemGroupConfig = {
       name: 'Make it real',
       icon: MakeItRealIcon,
       beta: true,
-      showWhen: () => true,
+      showWhen: notAllAIChatBlockShowWhen,
       handler: actionToHandler(
         'makeItReal',
         MakeItRealIconWithAnimation,

packages/frontend/core/src/blocksuite/presets/ai/index.ts

@@ -4,4 +4,5 @@ export { ChatPanel } from './chat-panel/index';
 export * from './entries/edgeless/actions-config';
 export * from './entries/index';
 export * from './messages/index';
+export { AIChatBlockPeekViewTemplate } from './peek-view/chat-block-peek-view';
 export * from './provider';

packages/frontend/core/src/blocksuite/presets/ai/messages/error.ts

@@ -1,6 +1,12 @@
 import { type EditorHost, WithDisposable } from '@blocksuite/block-std';
+import {
+  type AIError,
+  PaymentRequiredError,
+  UnauthorizedError,
+} from '@blocksuite/blocks';
 import { html, LitElement, nothing, type TemplateResult } from 'lit';
 import { customElement, property } from 'lit/decorators.js';
+import { styleMap } from 'lit/directives/style-map.js';
 
 import { ErrorTipIcon } from '../_common/icons';
 import { AIProvider } from '../provider';
@@ -112,3 +118,27 @@ declare global {
     'ai-error-wrapper': AIErrorWrapper;
   }
 }
+
+export function AIChatErrorRenderer(host: EditorHost, error: AIError) {
+  if (error instanceof PaymentRequiredError) {
+    return PaymentRequiredErrorRenderer(host);
+  } else if (error instanceof UnauthorizedError) {
+    return GeneralErrorRenderer(
+      html`You need to login to AFFiNE Cloud to continue using AFFiNE AI.`,
+      html`<div
+        style=${styleMap({
+          padding: '4px 12px',
+          borderRadius: '8px',
+          border: '1px solid var(--affine-border-color)',
+          cursor: 'pointer',
+          backgroundColor: 'var(--affine-hover-color)',
+        })}
+        @click=${() => AIProvider.slots.requestLogin.emit({ host })}
+      >
+        Login
+      </div>`
+    );
+  } else {
+    return GeneralErrorRenderer();
+  }
+}

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

@@ -0,0 +1,456 @@
+import type { EditorHost } from '@blocksuite/block-std';
+import { type AIError, openFileOrFiles } from '@blocksuite/blocks';
+import { type ChatMessage } from '@blocksuite/presets';
+import { css, html, LitElement, nothing } from 'lit';
+import { customElement, property, query, state } from 'lit/decorators.js';
+import { classMap } from 'lit/directives/class-map.js';
+import { repeat } from 'lit/directives/repeat.js';
+
+import {
+  ChatAbortIcon,
+  ChatClearIcon,
+  ChatSendIcon,
+  CloseIcon,
+  ImageIcon,
+} from '../_common/icons';
+import { AIProvider } from '../provider';
+import { reportResponse } from '../utils/action-reporter';
+import { readBlobAsURL } from '../utils/image';
+import type { ChatContext } from './types';
+
+const MaximumImageCount = 8;
+
+@customElement('chat-block-input')
+export class ChatBlockInput extends 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;
+      }
+      textarea::placeholder {
+        font-size: 14px;
+        font-weight: 400;
+        font-family: var(--affine-font-family);
+        color: var(--affine-placeholder-color);
+      }
+      textarea:focus {
+        outline: none;
+      }
+    }
+    .chat-input-images {
+      display: flex;
+      gap: 4px;
+      flex-wrap: wrap;
+      position: relative;
+      .image-container {
+        width: 58px;
+        height: 58px;
+        border-radius: 4px;
+        border: 1px solid var(--affine-border-color);
+        cursor: pointer;
+        overflow: hidden;
+        position: relative;
+        display: flex;
+        justify-content: center;
+        align-items: center;
+        img {
+          max-width: 100%;
+          max-height: 100%;
+          width: auto;
+          height: auto;
+        }
+      }
+    }
+    .close-wrapper {
+      width: 16px;
+      height: 16px;
+      border-radius: 4px;
+      border: 1px solid var(--affine-border-color);
+      justify-content: center;
+      align-items: center;
+      display: none;
+      position: absolute;
+      background-color: var(--affine-white);
+      z-index: 1;
+      cursor: pointer;
+    }
+    .close-wrapper:hover {
+      background-color: var(--affine-background-error-color);
+      border: 1px solid var(--affine-error-color);
+    }
+    .close-wrapper:hover svg path {
+      fill: var(--affine-error-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;
+      }
+    }
+
+    .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-send svg rect {
+          fill: ${this._isInputEmpty
+            ? 'var(--affine-text-disable-color)'
+            : 'var(--affine-primary-color)'};
+        }
+        .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;
+        }
+      </style>
+      <div class="ai-chat-input">
+        ${hasImages ? this._renderImages(images) : nothing}
+        <textarea
+          rows="1"
+          placeholder="What are your thoughts?"
+          @keydown=${async (evt: KeyboardEvent) => {
+            if (evt.key === 'Enter' && !evt.shiftKey && !evt.isComposing) {
+              evt.preventDefault();
+              await this._send();
+            }
+          }}
+          @input=${() => {
+            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';
+            }
+          }}
+          @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]);
+              }
+            }
+          }}
+        ></textarea>
+        <div class="chat-panel-input-actions">
+          <div
+            class=${cleanButtonClasses}
+            @click=${async () => {
+              if (disableCleanUp) {
+                return;
+              }
+              await this.cleanupHistories();
+            }}
+          >
+            ${ChatClearIcon}
+          </div>
+          ${images.length < MaximumImageCount
+            ? html`<div
+                class="image-upload"
+                @click=${async () => {
+                  const images = await openFileOrFiles({
+                    acceptType: 'Images',
+                    multiple: true,
+                  });
+                  if (!images) return;
+                  this._addImages(images);
+                }}
+              >
+                ${ImageIcon}
+              </div>`
+            : nothing}
+          ${status === 'transmitting'
+            ? html`<div
+                @click=${() => {
+                  this.chatContext.abortController?.abort();
+                  this.updateContext({ status: 'success' });
+                  reportResponse('aborted:stop');
+                }}
+              >
+                ${ChatAbortIcon}
+              </div>`
+            : html`<div @click="${this._send}" class="chat-panel-send">
+                ${ChatSendIcon}
+              </div>`}
+        </div>
+      </div>`;
+  }
+
+  @property({ attribute: false })
+  accessor parentSessionId!: string;
+
+  @property({ attribute: false })
+  accessor latestMessageId!: string;
+
+  @property({ attribute: false })
+  accessor host!: EditorHost;
+
+  @property({ attribute: false })
+  accessor updateChatBlock!: () => Promise<void>;
+
+  @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;
+
+  @query('.close-wrapper')
+  accessor closeWrapper: HTMLDivElement | null = null;
+
+  @state()
+  accessor _curIndex = -1;
+
+  private readonly _addImages = (images: File[]) => {
+    const oldImages = this.chatContext.images;
+    this.updateContext({
+      images: [...oldImages, ...images].slice(0, MaximumImageCount),
+    });
+  };
+
+  private _renderImages(images: File[]) {
+    return html`
+      <div
+        class="chat-input-images"
+        @mouseleave=${() => {
+          if (!this.closeWrapper) return;
+          this.closeWrapper.style.display = 'none';
+          this._curIndex = -1;
+        }}
+      >
+        ${repeat(
+          images,
+          image => image.name,
+          (image, index) =>
+            html`<div
+              class="image-container"
+              @mouseenter=${(evt: MouseEvent) => {
+                const ele = evt.target as HTMLImageElement;
+                const rect = ele.getBoundingClientRect();
+                if (!ele.parentElement) return;
+                const parentRect = ele.parentElement.getBoundingClientRect();
+                const left = Math.abs(rect.right - parentRect.left) - 8;
+                const top = Math.abs(parentRect.top - rect.top) - 8;
+                this._curIndex = index;
+                if (!this.closeWrapper) return;
+                this.closeWrapper.style.display = 'flex';
+                this.closeWrapper.style.left = left + 'px';
+                this.closeWrapper.style.top = top + 'px';
+              }}
+            >
+              <img src="${URL.createObjectURL(image)}" alt="${image.name}" />
+            </div>`
+        )}
+        <div
+          class="close-wrapper"
+          @click=${() => {
+            if (this._curIndex >= 0 && this._curIndex < images.length) {
+              const newImages = [...images];
+              newImages.splice(this._curIndex, 1);
+              this.updateContext({ images: newImages });
+              this._curIndex = -1;
+              if (!this.closeWrapper) return;
+              this.closeWrapper.style.display = 'none';
+            }
+          }}
+        >
+          ${CloseIcon}
+        </div>
+      </div>
+    `;
+  }
+
+  private readonly _send = async () => {
+    const { images, status } = this.chatContext;
+    if (status === 'loading' || status === 'transmitting') return;
+
+    const text = this.textarea.value;
+    if (!text && !images.length) {
+      return;
+    }
+
+    const { doc } = this.host;
+    this.textarea.value = '';
+    this._isInputEmpty = true;
+    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.chatContext.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(),
+        },
+      ],
+    });
+
+    const { currentChatBlockId, currentSessionId } = this.chatContext;
+    let content = '';
+    const chatBlockExists = !!currentChatBlockId;
+    try {
+      // If has not forked a chat session, fork a new one
+      let chatSessionId = currentSessionId;
+      if (!chatSessionId) {
+        const forkSessionId = await AIProvider.forkChat?.({
+          workspaceId: doc.collection.id,
+          docId: doc.id,
+          sessionId: this.parentSessionId,
+          latestMessageId: this.latestMessageId,
+        });
+        if (!forkSessionId) return;
+        this.updateContext({
+          currentSessionId: forkSessionId,
+        });
+        chatSessionId = forkSessionId;
+      }
+
+      const abortController = new AbortController();
+      const stream = AIProvider.actions.chat?.({
+        input: text,
+        sessionId: chatSessionId,
+        docId: doc.id,
+        attachments: images,
+        workspaceId: doc.collection.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.chatContext.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) {
+        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;
+  }
+}

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

@@ -0,0 +1,527 @@
+import './chat-block-input';
+import './date-time';
+import '../_common/components/chat-action-list';
+import '../_common/components/copy-more';
+
+import { type EditorHost } from '@blocksuite/block-std';
+import {
+  type AIError,
+  CanvasElementType,
+  ConnectorMode,
+  type EdgelessRootService,
+} from '@blocksuite/blocks';
+import { Bound } from '@blocksuite/global/utils';
+import {
+  type AIChatBlockModel,
+  type ChatMessage,
+  ChatMessagesSchema,
+} from '@blocksuite/presets';
+import type { Doc } from '@blocksuite/store';
+import { html, LitElement, nothing } from 'lit';
+import { customElement, property, query, state } from 'lit/decorators.js';
+import { classMap } from 'lit/directives/class-map.js';
+import { repeat } from 'lit/directives/repeat.js';
+
+import {
+  ChatBlockPeekViewActions,
+  constructUserInfoWithMessages,
+  queryHistoryMessages,
+} from '../_common/chat-actions-handle';
+import { SmallHintIcon } from '../_common/icons';
+import { AIChatErrorRenderer } from '../messages/error';
+import { AIProvider } from '../provider';
+import { PeekViewStyles } from './styles';
+import type { ChatContext } from './types';
+
+@customElement('ai-chat-block-peek-view')
+export class AIChatBlockPeekView extends LitElement {
+  static override styles = PeekViewStyles;
+
+  private get _rootService() {
+    return this.host.spec.getService('affine:page');
+  }
+
+  private get _modeService() {
+    return this._rootService.docModeService;
+  }
+
+  private get parentSessionId() {
+    return this.parentModel.sessionId;
+  }
+
+  private get historyMessagesString() {
+    return this.parentModel.messages;
+  }
+
+  private get parentXYWH() {
+    return this.parentModel.xywh;
+  }
+
+  private get parentChatBlockId() {
+    return this.parentModel.id;
+  }
+
+  private readonly _deserializeHistoryChatMessages = (
+    historyMessagesString: string
+  ) => {
+    try {
+      const result = ChatMessagesSchema.safeParse(
+        JSON.parse(historyMessagesString)
+      );
+      if (result.success) {
+        return result.data;
+      } else {
+        return [];
+      }
+    } catch {
+      return [];
+    }
+  };
+
+  private readonly _constructBranchChatBlockMessages = async (
+    doc: Doc,
+    forkSessionId: string
+  ) => {
+    const currentUserInfo = await AIProvider.userInfo;
+    const forkMessages = await queryHistoryMessages(doc, forkSessionId);
+    const forkLength = forkMessages.length;
+    const historyLength = this._historyMessages.length;
+
+    if (!forkLength || forkLength <= historyLength) {
+      return constructUserInfoWithMessages(forkMessages, currentUserInfo);
+    }
+
+    // Update history messages with the fork messages, keep user info
+    const historyMessages = this._historyMessages.map((message, idx) => {
+      return {
+        ...message,
+        id: forkMessages[idx]?.id ?? message.id,
+        attachments: [],
+      };
+    });
+
+    const currentChatMessages = constructUserInfoWithMessages(
+      forkMessages.slice(historyLength),
+      currentUserInfo
+    );
+    return [...historyMessages, ...currentChatMessages];
+  };
+
+  private readonly _resetContext = () => {
+    const { abortController } = this.chatContext;
+    if (abortController) {
+      abortController.abort();
+    }
+
+    this.updateContext({
+      status: 'idle',
+      error: null,
+      images: [],
+      abortController: null,
+      messages: [],
+      currentSessionId: null,
+      currentChatBlockId: null,
+    });
+  };
+
+  /**
+   * Create a new AI chat block based on the current session and history messages
+   */
+  createAIChatBlock = async () => {
+    // Only create AI chat block in edgeless mode
+    const mode = this._modeService.getMode();
+    if (mode !== 'edgeless') {
+      return;
+    }
+
+    // If there is already a chat block, do not create a new one
+    if (this.chatContext.currentChatBlockId) {
+      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
+    ) {
+      return;
+    }
+
+    const { doc } = this.host;
+    // create a new AI chat block
+    const surfaceBlock = doc
+      .getBlocks()
+      .find(block => block.flavour === 'affine:surface');
+    if (!surfaceBlock) {
+      return;
+    }
+
+    const parentXYWH = Bound.deserialize(this.parentXYWH);
+    const {
+      x: parentX,
+      y: parentY,
+      w: parentWidth,
+      h: parentHeight,
+    } = parentXYWH;
+
+    // Add AI chat block to the center of the viewport
+    // TODO: optimize the position of the AI chat block
+    const gap = parentWidth;
+    const x = parentX + parentWidth + gap;
+    const y = parentY;
+    const bound = new Bound(x, y, parentWidth, parentHeight);
+
+    // Get fork session messages
+    const messages = await this._constructBranchChatBlockMessages(
+      doc,
+      this.chatContext.currentSessionId
+    );
+    if (!messages.length) {
+      return;
+    }
+
+    const aiChatBlockId = doc.addBlock(
+      'affine:embed-ai-chat' as keyof BlockSuite.BlockModels,
+      {
+        xywh: bound.serialize(),
+        messages: JSON.stringify(messages),
+        sessionId: this.chatContext.currentSessionId,
+      },
+      surfaceBlock.id
+    );
+
+    if (!aiChatBlockId) {
+      return;
+    }
+
+    this.updateContext({ currentChatBlockId: aiChatBlockId });
+
+    // Connect the parent chat block to the AI chat block
+    const edgelessService = this._rootService as EdgelessRootService;
+    edgelessService.addElement(CanvasElementType.CONNECTOR, {
+      mode: ConnectorMode.Curve,
+      controllers: [],
+      source: { id: this.parentChatBlockId },
+      target: { id: aiChatBlockId },
+    });
+  };
+
+  /**
+   * Update the current chat messages with the new message
+   */
+  updateChatBlockMessages = async () => {
+    if (
+      !this.chatContext.currentChatBlockId ||
+      !this.chatContext.currentSessionId
+    ) {
+      return;
+    }
+
+    const { doc } = this.host;
+    const chatBlock = doc.getBlock(this.chatContext.currentChatBlockId);
+
+    // Get fork session messages
+    const messages = await this._constructBranchChatBlockMessages(
+      doc,
+      this.chatContext.currentSessionId
+    );
+    if (!messages.length) {
+      return;
+    }
+    doc.updateBlock(chatBlock.model, {
+      messages: JSON.stringify(messages),
+    });
+  };
+
+  updateContext = (context: Partial<ChatContext>) => {
+    this.chatContext = { ...this.chatContext, ...context };
+  };
+
+  /**
+   * Clean current chat messages and delete the newly created AI chat block
+   */
+  cleanCurrentChatHistories = async () => {
+    const { notificationService } = this._rootService;
+    if (!notificationService) {
+      return;
+    }
+
+    const { currentChatBlockId, currentSessionId } = this.chatContext;
+    if (!currentChatBlockId && !currentSessionId) {
+      return;
+    }
+
+    if (
+      await notificationService.confirm({
+        title: 'Clear History',
+        message:
+          'Are you sure you want to clear all history? This action will permanently delete all content, including all chat logs and data, and cannot be undone.',
+        confirmText: 'Confirm',
+        cancelText: 'Cancel',
+      })
+    ) {
+      const { doc } = this.host;
+      if (currentSessionId) {
+        await AIProvider.histories?.cleanup(doc.collection.id, doc.id, [
+          currentSessionId,
+        ]);
+      }
+
+      if (currentChatBlockId) {
+        const edgelessService = this._rootService as EdgelessRootService;
+        const chatBlock = doc.getBlock(currentChatBlockId).model;
+        if (chatBlock) {
+          const connectors = edgelessService.getConnectors(
+            chatBlock as AIChatBlockModel
+          );
+          doc.transact(() => {
+            // Delete the AI chat block
+            edgelessService.removeElement(currentChatBlockId);
+            // Delete the connectors
+            connectors.forEach(connector => {
+              edgelessService.removeElement(connector.id);
+            });
+          });
+        }
+      }
+
+      notificationService.toast('History cleared');
+      this._resetContext();
+    }
+  };
+
+  /**
+   * Retry the last chat message
+   */
+  retry = async () => {
+    const { doc } = this.host;
+    const { currentChatBlockId, currentSessionId } = this.chatContext;
+    if (!currentChatBlockId || !currentSessionId) {
+      return;
+    }
+
+    let content = '';
+    try {
+      const abortController = new AbortController();
+
+      const messages = [...this.chatContext.messages];
+      const last = messages[messages.length - 1];
+      if ('content' in last) {
+        last.content = '';
+        last.createdAt = new Date().toISOString();
+      }
+      this.updateContext({ messages, status: 'loading', error: null });
+
+      const stream = AIProvider.actions.chat?.({
+        sessionId: currentSessionId,
+        retry: true,
+        docId: doc.id,
+        workspaceId: doc.collection.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.chatContext.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) {
+      this.updateContext({ status: 'error', error: error as AIError });
+    } finally {
+      this.updateContext({ abortController: null });
+      if (content) {
+        // Update new chat block messages if there are contents returned from AI
+        await this.updateChatBlockMessages();
+      }
+    }
+  };
+
+  CurrentMessages = (currentMessages: ChatMessage[]) => {
+    if (!currentMessages.length) {
+      return nothing;
+    }
+
+    const { host } = this;
+    const actions = ChatBlockPeekViewActions;
+
+    return html`${repeat(
+      currentMessages,
+      message => message.createdAt + message.content,
+      (message, idx) => {
+        const { status, error } = this.chatContext;
+        const isAssistantMessage = message.role === 'assistant';
+        const isLastReply =
+          idx === currentMessages.length - 1 && isAssistantMessage;
+        const messageState =
+          isLastReply && status === 'transmitting' ? 'generating' : 'finished';
+        const shouldRenderError = isLastReply && status === 'error' && !!error;
+        const isNotReady = status === 'transmitting' || status === 'loading';
+        const shouldRenderCopyMore =
+          isAssistantMessage && !(isLastReply && isNotReady);
+        const shouldRenderActions =
+          isLastReply && !!message.content && !isNotReady;
+
+        const messageClasses = classMap({
+          'assistant-message-container': isAssistantMessage,
+        });
+
+        return html`<div class=${messageClasses}>
+          <ai-chat-message
+            .host=${host}
+            .message=${message}
+            .state=${messageState}
+          ></ai-chat-message>
+          ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
+          ${shouldRenderCopyMore
+            ? html` <chat-copy-more
+                .host=${host}
+                .actions=${actions}
+                .content=${message.content}
+                .isLast=${isLastReply}
+                .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
+                .messageId=${message.id ?? undefined}
+                .retry=${() => this.retry()}
+              ></chat-copy-more>`
+            : nothing}
+          ${shouldRenderActions
+            ? html`<chat-action-list
+                .host=${host}
+                .actions=${actions}
+                .content=${message.content}
+                .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
+                .messageId=${message.id ?? undefined}
+                .layoutDirection=${'horizontal'}
+              ></chat-action-list>`
+            : nothing}
+        </div>`;
+      }
+    )}`;
+  };
+
+  override connectedCallback() {
+    super.connectedCallback();
+    this._historyMessages = this._deserializeHistoryChatMessages(
+      this.historyMessagesString
+    );
+    queryHistoryMessages(this.host.doc, this.parentSessionId)
+      .then(messages => {
+        this._historyMessages = this._historyMessages.map((message, idx) => {
+          return {
+            ...message,
+            attachments: messages[idx]?.attachments ?? [],
+          };
+        });
+      })
+      .catch((err: Error) => {
+        console.error('Query history messages failed', err);
+      });
+  }
+
+  override firstUpdated() {
+    // first time render, scroll ai-chat-messages-container to bottom
+    requestAnimationFrame(() => {
+      if (this._chatMessagesContainer) {
+        this._chatMessagesContainer.scrollTop =
+          this._chatMessagesContainer.scrollHeight;
+      }
+    });
+  }
+
+  override render() {
+    const { host, _historyMessages } = this;
+    if (!_historyMessages.length) {
+      return nothing;
+    }
+
+    const latestHistoryMessage = _historyMessages[_historyMessages.length - 1];
+    const latestMessageCreatedAt = latestHistoryMessage.createdAt;
+    const latestHistoryMessageId = latestHistoryMessage.id;
+    const {
+      parentSessionId,
+      updateChatBlockMessages,
+      createAIChatBlock,
+      cleanCurrentChatHistories,
+      chatContext,
+      updateContext,
+    } = this;
+
+    const { messages: currentChatMessages } = chatContext;
+
+    return html`<div class="ai-chat-block-peek-view-container">
+      <div class="ai-chat-messages-container">
+        <ai-chat-messages
+          .host=${host}
+          .messages=${_historyMessages}
+        ></ai-chat-messages>
+        <date-time .date=${latestMessageCreatedAt}></date-time>
+        <div class="new-chat-messages-container">
+          ${this.CurrentMessages(currentChatMessages)}
+        </div>
+      </div>
+      <chat-block-input
+        .host=${host}
+        .parentSessionId=${parentSessionId}
+        .latestMessageId=${latestHistoryMessageId}
+        .updateChatBlock=${updateChatBlockMessages}
+        .createChatBlock=${createAIChatBlock}
+        .cleanupHistories=${cleanCurrentChatHistories}
+        .chatContext=${chatContext}
+        .updateContext=${updateContext}
+      ></chat-block-input>
+      <div class="peek-view-footer">
+        ${SmallHintIcon}
+        <div>AI outputs can be misleading or wrong</div>
+      </div>
+    </div> `;
+  }
+
+  @query('.ai-chat-messages-container')
+  accessor _chatMessagesContainer!: HTMLDivElement;
+
+  @property({ attribute: false })
+  accessor parentModel!: AIChatBlockModel;
+
+  @property({ attribute: false })
+  accessor host!: EditorHost;
+
+  @state()
+  accessor _historyMessages: ChatMessage[] = [];
+
+  @state()
+  accessor chatContext: ChatContext = {
+    status: 'idle',
+    error: null,
+    images: [],
+    abortController: null,
+    messages: [],
+    currentSessionId: null,
+    currentChatBlockId: null,
+  };
+}
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'ai-chat-block-peek-view': AIChatBlockPeekView;
+  }
+}
+
+export const AIChatBlockPeekViewTemplate = (
+  parentModel: AIChatBlockModel,
+  host: EditorHost
+) => {
+  return html`<ai-chat-block-peek-view
+    .parentModel=${parentModel}
+    .host=${host}
+  ></ai-chat-block-peek-view>`;
+};

packages/frontend/core/src/blocksuite/presets/ai/peek-view/date-time.ts

@@ -0,0 +1,59 @@
+import { i18nTime } from '@affine/i18n';
+import { css, html, LitElement } from 'lit';
+import { customElement, property } from 'lit/decorators.js';
+
+@customElement('date-time')
+export class DateTime extends LitElement {
+  static override styles = css`
+    :host {
+      width: 100%;
+    }
+    .date-time-container {
+      width: 100%;
+      display: flex;
+      justify-content: space-between;
+      align-items: center;
+    }
+    .line {
+      flex-grow: 1;
+      height: 0.5px;
+      background-color: var(--affine-border-color);
+    }
+    .date-time {
+      padding: 0 8px;
+      font-size: var(--affine-font-xs);
+      font-weight: 400;
+      line-height: 22px;
+      text-align: center;
+      color: var(--affine-text-secondary-color);
+    }
+  `;
+
+  override render() {
+    const date = i18nTime(this.date, {
+      relative: {
+        max: [1, 'day'],
+        accuracy: 'minute',
+        weekday: true,
+      },
+      absolute: {
+        accuracy: 'second',
+      },
+    });
+
+    return html`<div class="date-time-container">
+      <div class="line"></div>
+      <div class="date-time">${date}</div>
+      <div class="line"></div>
+    </div>`;
+  }
+
+  @property({ attribute: false })
+  accessor date!: string;
+}
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'date-time': DateTime;
+  }
+}

packages/frontend/core/src/blocksuite/presets/ai/peek-view/styles.ts

@@ -0,0 +1,73 @@
+import { baseTheme } from '@toeverything/theme';
+import { css, unsafeCSS } from 'lit';
+
+export const PeekViewStyles = css`
+  :host {
+    width: 100%;
+    height: 100%;
+  }
+
+  .ai-chat-block-peek-view-container {
+    gap: 8px;
+    width: 100%;
+    height: 100%;
+    display: flex;
+    align-items: center;
+    box-sizing: border-box;
+    justify-content: start;
+    flex-direction: column;
+    box-sizing: border-box;
+    padding: 24px 120px 16px 120px;
+    font-family: ${unsafeCSS(baseTheme.fontSansFamily)};
+  }
+
+  .ai-chat-messages-container {
+    display: flex;
+    align-items: center;
+    justify-content: start;
+    flex-direction: column;
+    box-sizing: border-box;
+    width: 100%;
+    color: var(--affine-text-primary-color);
+    line-height: 22px;
+    font-size: var(--affine-font-sm);
+    overflow-y: auto;
+    overflow-x: hidden;
+    flex: 1;
+    gap: 24px;
+    -ms-overflow-style: none; /* IE and Edge */
+    scrollbar-width: none; /* Firefox */
+  }
+
+  .new-chat-messages-container {
+    width: 100%;
+    box-sizing: border-box;
+    min-height: 450px;
+    display: flex;
+    flex-direction: column;
+    gap: 24px;
+  }
+
+  .ai-chat-messages-container::-webkit-scrollbar {
+    display: none;
+  }
+
+  .assistant-message-container {
+    width: 100%;
+    box-sizing: border-box;
+    display: flex;
+    flex-direction: column;
+    gap: 8px;
+  }
+
+  .peek-view-footer {
+    padding: 0 12px;
+    width: 100%;
+    height: 20px;
+    display: flex;
+    gap: 4px;
+    align-items: center;
+    color: var(--affine-text-secondary-color);
+    font-size: var(--affine-font-xs);
+  }
+`;

packages/frontend/core/src/blocksuite/presets/ai/peek-view/types.ts

@@ -0,0 +1,19 @@
+import type { AIError } from '@blocksuite/blocks';
+import { type ChatMessage } from '@blocksuite/presets';
+
+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;
+};

packages/frontend/core/src/blocksuite/presets/ai/provider.ts

@@ -78,6 +78,8 @@ export class AIProvider {
 
   static LAST_ACTION_SESSIONID = '';
 
+  static LAST_ROOT_SESSION_ID = '';
+
   static MAX_LOCAL_HISTORY = 10;
 
   private readonly actions: Partial<BlockSuitePresets.AIActions> = {};

packages/frontend/core/src/blocksuite/presets/ai/utils/editor-actions.ts

@@ -59,9 +59,11 @@ export const insert = async (
   );
   const insertIndex = below ? index + 1 : index;
 
+  const { doc } = host;
   const models = await insertFromMarkdown(
     host,
     content,
+    doc,
     blockParent.model.id,
     insertIndex
   );
@@ -110,9 +112,11 @@ export const replace = async (
       host.doc.deleteBlock(model);
     });
 
+    const { doc } = host;
     const models = await insertFromMarkdown(
       host,
       content,
+      doc,
       firstBlockParent.model.id,
       firstIndex
     );

packages/frontend/core/src/blocksuite/presets/ai/utils/markdown-utils.ts

@@ -148,6 +148,7 @@ export const markdownToSnapshot = async (
 export async function insertFromMarkdown(
   host: EditorHost,
   markdown: string,
+  doc: Doc,
   parent?: string,
   index?: number
 ) {
@@ -160,7 +161,7 @@ export async function insertFromMarkdown(
     const blockSnapshot = snapshots[i];
     const model = await job.snapshotToBlock(
       blockSnapshot,
-      host.std.doc,
+      doc,
       parent,
       (index ?? 0) + i
     );

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts

@@ -24,6 +24,7 @@ export type TextToTextOptions = {
   signal?: AbortSignal;
   retry?: boolean;
   workflow?: boolean;
+  isRootSession?: boolean;
   postfix?: (text: string) => string;
 };
 
@@ -151,6 +152,7 @@ export function textToText({
   timeout = TIMEOUT,
   retry = false,
   workflow = false,
+  isRootSession = false,
   postfix,
 }: TextToTextOptions) {
   let _sessionId: string;
@@ -188,6 +190,9 @@ export function textToText({
           workflow ? 'workflow' : undefined
         );
         AIProvider.LAST_ACTION_SESSIONID = _sessionId;
+        if (isRootSession) {
+          AIProvider.LAST_ROOT_SESSION_ID = _sessionId;
+        }
 
         if (signal) {
           if (signal.aborted) {
@@ -250,6 +255,10 @@ export function textToText({
         }
 
         AIProvider.LAST_ACTION_SESSIONID = _sessionId;
+        if (isRootSession) {
+          AIProvider.LAST_ROOT_SESSION_ID = _sessionId;
+        }
+
         return client.chatText({
           sessionId: _sessionId,
           messageId: _messageId,

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx

@@ -77,7 +77,8 @@ function setupAIProvider() {
 
   //#region actions
   AIProvider.provide('chat', options => {
-    const sessionId = getChatSessionId(options.workspaceId, options.docId);
+    const sessionId =
+      options.sessionId ?? getChatSessionId(options.workspaceId, options.docId);
     return textToText({
       ...options,
       content: options.input,

packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/tracker.ts

@@ -18,13 +18,15 @@ type AIActionEventProperties = {
     | 'AI action panel'
     | 'right side bar'
     | 'inline chat panel'
-    | 'AI result panel';
+    | 'AI result panel'
+    | 'AI chat block';
   module:
     | 'exit confirmation'
     | 'AI action panel'
     | 'AI chat panel'
     | 'inline chat panel'
-    | 'AI result panel';
+    | 'AI result panel'
+    | 'AI chat block';
   control:
     | 'stop button'
     | 'format toolbar'
@@ -139,6 +141,8 @@ function inferSegment(
     return 'AI result panel';
   } else if (event.options.where === 'chat-panel') {
     return 'right side bar';
+  } else if (event.options.where === 'ai-chat-block') {
+    return 'AI chat block';
   } else {
     return 'AI action panel';
   }
@@ -155,6 +159,8 @@ function inferModule(
     return 'AI result panel';
   } else if (event.options.where === 'inline-chat-panel') {
     return 'inline chat panel';
+  } else if (event.options.where === 'ai-chat-block') {
+    return 'AI chat block';
   } else {
     return 'AI action panel';
   }

packages/frontend/core/src/modules/peek-view/entities/peek-view.ts

@@ -1,4 +1,4 @@
-import type { BlockComponent } from '@blocksuite/block-std';
+import type { BlockComponent, EditorHost } from '@blocksuite/block-std';
 import {
   AffineReference,
   type EmbedLinkedDocModel,
@@ -7,6 +7,7 @@ import {
   type SurfaceRefBlockComponent,
   type SurfaceRefBlockModel,
 } from '@blocksuite/blocks';
+import type { AIChatBlockModel } from '@blocksuite/presets';
 import type { BlockModel } from '@blocksuite/store';
 import { type DocMode, Entity, LiveData } from '@toeverything/infra';
 import type { TemplateResult } from 'lit';
@@ -36,6 +37,13 @@ export type ImagePeekViewInfo = {
   blockId: string;
 };
 
+export type AIChatBlockPeekViewInfo = {
+  type: 'ai-chat-block';
+  docId: string;
+  host: EditorHost;
+  model: AIChatBlockModel;
+};
+
 export type CustomTemplatePeekViewInfo = {
   type: 'template';
   template: TemplateResult;
@@ -43,7 +51,11 @@ export type CustomTemplatePeekViewInfo = {
 
 export type ActivePeekView = {
   target: PeekViewTarget;
-  info: DocPeekViewInfo | ImagePeekViewInfo | CustomTemplatePeekViewInfo;
+  info:
+    | DocPeekViewInfo
+    | ImagePeekViewInfo
+    | CustomTemplatePeekViewInfo
+    | AIChatBlockPeekViewInfo;
 };
 
 const EMBED_DOC_FLAVOURS = [
@@ -69,6 +81,12 @@ const isSurfaceRefModel = (
   return blockModel.flavour === 'affine:surface-ref';
 };
 
+const isAIChatBlockModel = (
+  blockModel: BlockModel
+): blockModel is AIChatBlockModel => {
+  return blockModel.flavour === 'affine:embed-ai-chat';
+};
+
 function resolvePeekInfoFromPeekTarget(
   peekTarget: PeekViewTarget,
   template?: TemplateResult
@@ -113,6 +131,13 @@ function resolvePeekInfoFromPeekTarget(
         docId: blockModel.doc.id,
         blockId: blockModel.id,
       };
+    } else if (isAIChatBlockModel(blockModel)) {
+      return {
+        type: 'ai-chat-block',
+        docId: blockModel.doc.id,
+        model: blockModel,
+        host: peekTarget.host,
+      };
     }
   } else if (peekTarget instanceof HTMLAnchorElement) {
     const maybeDoc = resolveLinkToDoc(peekTarget.href);

packages/frontend/core/src/modules/peek-view/view/peek-view-manager.tsx

@@ -1,4 +1,5 @@
 import { toReactNode } from '@affine/component';
+import { AIChatBlockPeekViewTemplate } from '@affine/core/blocksuite/presets/ai';
 import { BlockComponent } from '@blocksuite/block-std';
 import { useLiveData, useService } from '@toeverything/infra';
 import { useEffect, useMemo } from 'react';
@@ -35,6 +36,11 @@ function renderPeekView({ info }: ActivePeekView) {
     return <ImagePreviewPeekView docId={info.docId} blockId={info.blockId} />;
   }
 
+  if (info.type === 'ai-chat-block') {
+    const template = AIChatBlockPeekViewTemplate(info.model, info.host);
+    return toReactNode(template);
+  }
+
   return null; // unreachable
 }