feat(core): after the user clicks ask ai, the input pops up directly (#9039)

Issue [AF-1762](https://linear.app/affine-design/issue/AF-1762).

Ask AI on edgeless mode:

<div class='graphite__hidden'>
          <div>🎥 Video uploaded on Graphite:</div>
            <a href="https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/1218afba-466f-4570-afd4-679b6b09cc8d.mov">
              <img src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/1218afba-466f-4570-afd4-679b6b09cc8d.mov">
            </a>
          </div>
<video src="https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/sJGviKxfE3Ap685cl5bj/1218afba-466f-4570-afd4-679b6b09cc8d.mov">录屏2024-12-06 09.54.52.mov</video>

Ask AI on page mode:
<div class='graphite__hidden'>
          <div>🎥 Video uploaded on Graphite:</div>
            <a href="https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/58d19705-f646-4957-8628-15845b47222b.mov">
              <img src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/58d19705-f646-4957-8628-15845b47222b.mov">
            </a>
          </div>
<video src="https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/sJGviKxfE3Ap685cl5bj/58d19705-f646-4957-8628-15845b47222b.mov">录屏2024-12-06 09.52.51.mov</video>

packages/frontend/core/src/blocksuite/presets/ai/_common/components/ask-ai-button.ts

@@ -3,11 +3,10 @@ import './ask-ai-panel';
 import { type EditorHost } from '@blocksuite/affine/block-std';
 import {
   type AIItemGroupConfig,
-  AIStarIcon,
   EdgelessRootService,
 } from '@blocksuite/affine/blocks';
 import { createLitPortal, HoverController } from '@blocksuite/affine/blocks';
-import { assertExists, WithDisposable } from '@blocksuite/affine/global/utils';
+import { WithDisposable } from '@blocksuite/affine/global/utils';
 import { flip, offset } from '@floating-ui/dom';
 import { css, html, LitElement, nothing } from 'lit';
 import { property, query } from 'lit/decorators.js';
@@ -15,24 +14,12 @@ import { ref } from 'lit/directives/ref.js';
 import { styleMap } from 'lit/directives/style-map.js';
 
 import { getRootService } from '../../utils/selection-utils';
+import type { ButtonSize } from './ask-ai-icon';
 
-type buttonSize = 'small' | 'middle' | 'large';
 type toggleType = 'hover' | 'click';
 
-const buttonWidthMap: Record<buttonSize, string> = {
-  small: '72px',
-  middle: '76px',
-  large: '82px',
-};
-
-const buttonHeightMap: Record<buttonSize, string> = {
-  small: '24px',
-  middle: '32px',
-  large: '32px',
-};
-
 export type AskAIButtonOptions = {
-  size: buttonSize;
+  size: ButtonSize;
   backgroundColor?: string;
   boxShadow?: string;
   panelWidth?: number;
@@ -53,39 +40,6 @@ export class AskAIButton extends WithDisposable(LitElement) {
       position: relative;
       user-select: none;
     }
-
-    .ask-ai-icon-button {
-      display: flex;
-      align-items: center;
-      justify-content: center;
-      color: var(--affine-brand-color);
-      font-size: var(--affine-font-sm);
-      font-weight: 500;
-    }
-
-    .ask-ai-icon-button.small {
-      font-size: var(--affine-font-xs);
-      svg {
-        scale: 0.8;
-        margin-right: 2px;
-      }
-    }
-
-    .ask-ai-icon-button.large {
-      font-size: var(--affine-font-md);
-      svg {
-        scale: 1.2;
-      }
-    }
-
-    .ask-ai-icon-button span {
-      line-height: 22px;
-    }
-
-    .ask-ai-icon-button svg {
-      margin-right: 4px;
-      color: var(--affine-brand-color);
-    }
   `;
 
   @query('.ask-ai-button')
@@ -127,7 +81,7 @@ export class AskAIButton extends WithDisposable(LitElement) {
 
   @property({ attribute: false })
   accessor options: AskAIButtonOptions = {
-    size: 'middle',
+    size: 'small',
     backgroundColor: undefined,
     boxShadow: undefined,
     panelWidth: 330,
@@ -145,13 +99,16 @@ export class AskAIButton extends WithDisposable(LitElement) {
       return;
     }
 
+    if (!this._askAIButton) {
+      return;
+    }
+
     if (this._abortController) {
       this._clearAbortController();
       return;
     }
 
     this._abortController = new AbortController();
-    assertExists(this._askAIButton);
     const panelMinWidth = this.options.panelWidth || 330;
     createLitPortal({
       template: html`<ask-ai-panel
@@ -177,7 +134,7 @@ export class AskAIButton extends WithDisposable(LitElement) {
   }
 
   override render() {
-    const { size = 'small', backgroundColor, boxShadow } = this.options;
+    const { size, backgroundColor, boxShadow } = this.options;
     const { toggleType } = this;
     const buttonStyles = styleMap({
       backgroundColor: backgroundColor || 'transparent',
@@ -189,13 +146,7 @@ export class AskAIButton extends WithDisposable(LitElement) {
       ${toggleType === 'hover' ? ref(this._whenHover.setReference) : nothing}
       @click=${this._toggleAIPanel}
     >
-      <icon-button
-        class="ask-ai-icon-button ${size}"
-        width=${buttonWidthMap[size]}
-        height=${buttonHeightMap[size]}
-      >
-        ${AIStarIcon} <span>Ask AI</span></icon-button
-      >
+      <ask-ai-icon .size=${size}></ask-ai-icon>
     </div>`;
   }
 }

packages/frontend/core/src/blocksuite/presets/ai/_common/components/ask-ai-icon.ts

@@ -0,0 +1,77 @@
+import { AIStarIcon } from '@blocksuite/affine/blocks';
+import { WithDisposable } from '@blocksuite/affine/global/utils';
+import { css, html, LitElement } from 'lit';
+import { property } from 'lit/decorators.js';
+
+export type ButtonSize = 'small' | 'middle' | 'large';
+
+const buttonWidthMap: Record<ButtonSize, string> = {
+  small: '72px',
+  middle: '76px',
+  large: '82px',
+};
+
+const buttonHeightMap: Record<ButtonSize, string> = {
+  small: '24px',
+  middle: '32px',
+  large: '32px',
+};
+
+export class AskAIIcon extends WithDisposable(LitElement) {
+  @property({ attribute: false })
+  accessor size!: ButtonSize;
+
+  static override styles = css`
+    .ask-ai-icon-button {
+      display: flex;
+      align-items: center;
+      justify-content: center;
+      color: var(--affine-brand-color);
+      font-size: var(--affine-font-sm);
+      font-weight: 500;
+    }
+
+    .ask-ai-icon-button.small {
+      font-size: var(--affine-font-xs);
+      svg {
+        scale: 0.8;
+        margin-right: 2px;
+      }
+    }
+
+    .ask-ai-icon-button.large {
+      font-size: var(--affine-font-md);
+      svg {
+        scale: 1.2;
+      }
+    }
+
+    .ask-ai-icon-button span {
+      line-height: 22px;
+    }
+
+    .ask-ai-icon-button svg {
+      margin-right: 4px;
+      color: var(--affine-brand-color);
+    }
+  `;
+
+  override render() {
+    return html`
+      <icon-button
+        class="ask-ai-icon-button ${this.size}"
+        width=${buttonWidthMap[this.size]}
+        height=${buttonHeightMap[this.size]}
+      >
+        ${AIStarIcon}
+        <span>Ask AI</span>
+      </icon-button>
+    `;
+  }
+}
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'ask-ai-icon': AskAIIcon;
+  }
+}

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

@@ -44,6 +44,9 @@ export class AskAIPanel extends WithDisposable(LitElement) {
   @property({ attribute: false })
   accessor abortController: AbortController | null = null;
 
+  @property({ attribute: false })
+  accessor onItemClick: (() => void) | undefined = undefined;
+
   @property({ attribute: false })
   accessor minWidth = 330;
 
@@ -81,6 +84,7 @@ export class AskAIPanel extends WithDisposable(LitElement) {
       <ai-item-list
         .host=${this.host}
         .groups=${this._actionGroups}
+        .onClick=${this.onItemClick}
       ></ai-item-list>
     </div>`;
   }

packages/frontend/core/src/blocksuite/presets/ai/_common/components/ask-ai-toolbar.ts

@@ -0,0 +1,118 @@
+import type { EditorHost } from '@blocksuite/affine/block-std';
+import {
+  type AffineAIPanelWidgetConfig,
+  type AIItemGroupConfig,
+  createLitPortal,
+} from '@blocksuite/affine/blocks';
+import { WithDisposable } from '@blocksuite/affine/global/utils';
+import { flip, offset } from '@floating-ui/dom';
+import { css, html, LitElement } from 'lit';
+import { property } from 'lit/decorators.js';
+
+import { getAIPanel } from '../../ai-panel';
+import { AIProvider } from '../../provider';
+import { extractContext } from '../../utils/extract';
+
+export class AskAIToolbarButton extends WithDisposable(LitElement) {
+  static override styles = css`
+    .ask-ai-button {
+      border-radius: 4px;
+      position: relative;
+      user-select: none;
+    }
+  `;
+
+  private _abortController: AbortController | null = null;
+
+  private _panelRoot: HTMLDivElement | null = null;
+
+  @property({ attribute: false })
+  accessor host!: EditorHost;
+
+  @property({ attribute: false })
+  accessor actionGroups!: AIItemGroupConfig[];
+
+  private readonly _onItemClick = () => {
+    const aiPanel = getAIPanel(this.host);
+    aiPanel.restoreSelection();
+    this._clearAbortController();
+  };
+
+  private readonly _clearAbortController = () => {
+    this._abortController?.abort();
+    this._abortController = null;
+  };
+
+  private readonly _openAIPanel = () => {
+    this._clearAbortController();
+    const aiPanel = getAIPanel(this.host);
+    this._abortController = new AbortController();
+    this._panelRoot = createLitPortal({
+      template: html`
+        <ask-ai-panel
+          .host=${this.host}
+          .actionGroups=${this.actionGroups}
+          .onItemClick=${this._onItemClick}
+        ></ask-ai-panel>
+      `,
+      computePosition: {
+        referenceElement: aiPanel,
+        placement: 'top-start',
+        middleware: [flip(), offset({ mainAxis: 3 })],
+        autoUpdate: true,
+      },
+      abortController: this._abortController,
+      closeOnClickAway: true,
+    });
+  };
+
+  private readonly _generateAnswer: AffineAIPanelWidgetConfig['generateAnswer'] =
+    ({ finish, input }) => {
+      finish('success');
+      const aiPanel = getAIPanel(this.host);
+      aiPanel.discard();
+      AIProvider.slots.requestOpenWithChat.emit({ host: this.host });
+      extractContext(this.host)
+        .then(context => {
+          AIProvider.slots.requestSendWithChat.emit({ input, context });
+        })
+        .catch(console.error);
+    };
+
+  private readonly _onClick = () => {
+    const aiPanel = getAIPanel(this.host);
+    if (!aiPanel.config) return;
+    aiPanel.config.generateAnswer = this._generateAnswer;
+    aiPanel.config.inputCallback = text => {
+      if (!this._panelRoot) return;
+      this._panelRoot.style.visibility = text ? 'hidden' : 'visible';
+    };
+
+    const textSelection = this.host.selection.find('text');
+    const blockSelections = this.host.selection.filter('block');
+    let lastBlockId: string | undefined;
+    if (textSelection) {
+      lastBlockId = textSelection.to?.blockId ?? textSelection.blockId;
+    } else if (blockSelections.length) {
+      lastBlockId = blockSelections[blockSelections.length - 1].blockId;
+    }
+    if (!lastBlockId) return;
+    const block = this.host.view.getBlock(lastBlockId);
+    if (!block) return;
+    aiPanel.setState('input', block);
+
+    setTimeout(() => this._openAIPanel(), 0);
+  };
+
+  override render() {
+    return html`<div class="ask-ai-button" @click=${this._onClick}>
+      <ask-ai-icon .size=${'middle'}></ask-ai-icon>
+    </div>`;
+  }
+}
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'ask-ai-toolbar-button': AskAIToolbarButton;
+  }
+}

packages/frontend/core/src/blocksuite/presets/ai/_common/config.ts

@@ -399,7 +399,11 @@ const OthersAIGroup: AIItemGroupConfig = {
       icon: CommentIcon,
       handler: host => {
         const panel = getAIPanel(host);
-        AIProvider.slots.requestOpenWithChat.emit({ host, autoSelect: true });
+        AIProvider.slots.requestOpenWithChat.emit({
+          host,
+          autoSelect: true,
+          appendCard: true,
+        });
         panel.hide();
       },
     },

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

@@ -506,7 +506,9 @@ export class ChatCards extends WithDisposable(LitElement) {
 
     this._disposables.add(
       AIProvider.slots.requestOpenWithChat.on(async params => {
-        await this._appendCardWithParams(params);
+        if (params.appendCard) {
+          await this._appendCardWithParams(params);
+        }
       })
     );
 

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

@@ -296,6 +296,18 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
     `;
   }
 
+  override connectedCallback() {
+    super.connectedCallback();
+
+    this._disposables.add(
+      AIProvider.slots.requestSendWithChat.on(async ({ input, context }) => {
+        context && this.updateContext(context);
+        await this.updateComplete;
+        await this.send(input);
+      })
+    );
+  }
+
   protected override render() {
     const { images, status } = this.chatContextValue;
     const hasImages = images.length > 0;
@@ -420,11 +432,11 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
       </div>`;
   }
 
-  send = async () => {
+  send = async (input?: string) => {
     const { status, markdown } = this.chatContextValue;
     if (status === 'loading' || status === 'transmitting') return;
 
-    const text = this.textarea.value;
+    const text = input || this.textarea.value;
     const { images } = this.chatContextValue;
     if (!text && images.length === 0) {
       return;

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

@@ -263,7 +263,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
         ></chat-cards>
       </div>
       ${this.showDownIndicator
-        ? html`<div class="down-indicator" @click=${() => this.scrollToDown()}>
+        ? html`<div class="down-indicator" @click=${this.scrollToEnd}>
             ${DownArrowIcon}
           </div>`
         : nothing} `;
@@ -387,8 +387,15 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
     return html` <ai-loading></ai-loading>`;
   }
 
-  scrollToDown() {
-    this.messagesContainer.scrollTo(0, this.messagesContainer.scrollHeight);
+  scrollToEnd() {
+    this.updateComplete
+      .then(() => {
+        this.messagesContainer.scrollTo({
+          top: this.messagesContainer.scrollHeight,
+          behavior: 'smooth',
+        });
+      })
+      .catch(console.error);
   }
 
   retry = async () => {

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

@@ -133,7 +133,7 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
       };
 
       this.isLoading = false;
-      this.scrollToDown();
+      this._scrollToEnd();
     })().catch(console.error);
   }, 200);
 
@@ -158,6 +158,10 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
     chatSessionId: null,
   };
 
+  private readonly _scrollToEnd = () => {
+    requestAnimationFrame(() => this._chatMessages.value?.scrollToEnd());
+  };
+
   private readonly _cleanupHistories = async () => {
     const notification = this.host.std.getOptional(NotificationProvider);
     if (!notification) return;
@@ -190,6 +194,19 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
       this.chatContextValue.chatSessionId = null;
       this._resetItems();
     }
+
+    if (!this.isLoading && _changedProperties.has('chatContextValue')) {
+      if (this.chatContextValue.status !== 'idle') {
+        this._scrollToEnd();
+      }
+      if (
+        this.chatContextValue.status === 'loading' ||
+        this.chatContextValue.status === 'error' ||
+        this.chatContextValue.status === 'success'
+      ) {
+        setTimeout(this._scrollToEnd, 500);
+      }
+    }
   }
 
   override connectedCallback() {
@@ -237,10 +254,6 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
     });
   };
 
-  scrollToDown() {
-    requestAnimationFrame(() => this._chatMessages.value?.scrollToDown());
-  }
-
   override render() {
     return html` <div class="chat-panel-container">
       <div class="chat-panel-title">

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

@@ -110,6 +110,7 @@ const othersGroup: AIItemGroupConfig = {
           host,
           mode: 'edgeless',
           autoSelect: true,
+          appendCard: true,
         });
         panel.hide();
       },

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

@@ -9,6 +9,10 @@ import { EdgelessCopilotToolbarEntry } from '@blocksuite/affine/blocks';
 import { noop } from '@blocksuite/affine/global/utils';
 import { html } from 'lit';
 
+import { getAIPanel } from '../../ai-panel';
+import { AIProvider } from '../../provider';
+import { getEdgelessCopilotWidget } from '../../utils/edgeless';
+import { extractContext } from '../../utils/extract';
 import { edgelessActionGroups } from './actions-config';
 
 noop(EdgelessCopilotToolbarEntry);
@@ -38,10 +42,36 @@ export function setupEdgelessElementToolbarAIEntry(
 
       if (filteredGroups.every(group => group.items.length === 0)) return null;
 
+      const handler = () => {
+        const aiPanel = getAIPanel(edgeless.host);
+        if (aiPanel.config) {
+          aiPanel.config.generateAnswer = ({ finish, input }) => {
+            finish('success');
+            aiPanel.discard();
+            AIProvider.slots.requestOpenWithChat.emit({ host: edgeless.host });
+            extractContext(edgeless.host)
+              .then(context => {
+                AIProvider.slots.requestSendWithChat.emit({ input, context });
+              })
+              .catch(console.error);
+          };
+          aiPanel.config.inputCallback = text => {
+            const copilotWidget = getEdgelessCopilotWidget(edgeless.host);
+            const panel = copilotWidget.shadowRoot?.querySelector(
+              'edgeless-copilot-panel'
+            );
+            if (panel instanceof HTMLElement) {
+              panel.style.visibility = text ? 'hidden' : 'visible';
+            }
+          };
+        }
+      };
+
       return html`<edgeless-copilot-toolbar-entry
         .edgeless=${edgeless}
         .host=${edgeless.host}
         .groups=${edgelessActionGroups}
+        .onClick=${handler}
       ></edgeless-copilot-toolbar-entry>`;
     },
   });

packages/frontend/core/src/blocksuite/presets/ai/entries/format-bar/setup-format-bar.ts

@@ -15,11 +15,12 @@ export function setupFormatBarAIEntry(formatBar: AffineFormatBarWidget) {
       {
         type: 'custom' as const,
         render(formatBar: AffineFormatBarWidget): TemplateResult | null {
-          return html` <ask-ai-button
-            .host=${formatBar.host}
-            .actionGroups=${AIItemGroups}
-            .toggleType=${'hover'}
-          ></ask-ai-button>`;
+          return html`
+            <ask-ai-toolbar-button
+              .host=${formatBar.host}
+              .actionGroups=${AIItemGroups}
+            ></ask-ai-toolbar-button>
+          `;
         },
       },
       { type: 'divider' },

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

@@ -131,6 +131,7 @@ declare global {
 }
 
 export function AIChatErrorRenderer(host: EditorHost, error: AIError) {
+  console.error(error);
   if (error instanceof PaymentRequiredError) {
     return PaymentRequiredErrorRenderer(host);
   } else if (error instanceof UnauthorizedError) {

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

@@ -6,6 +6,8 @@ import {
 import { Slot } from '@blocksuite/affine/store';
 import { captureException } from '@sentry/react';
 
+import type { ChatContextValue } from './chat-panel/chat-context';
+
 export interface AIUserInfo {
   id: string;
   email: string;
@@ -18,6 +20,12 @@ export interface AIChatParams {
   mode?: 'page' | 'edgeless';
   // Auto select and append selection to input via `Continue with AI` action.
   autoSelect?: boolean;
+  appendCard?: boolean;
+}
+
+export interface AISendParams {
+  input?: string;
+  context?: Partial<ChatContextValue | null>;
 }
 
 export type ActionEventType =
@@ -104,6 +112,7 @@ export class AIProvider {
     // use case: when user selects "continue in chat" in an ask ai result panel
     // do we need to pass the context to the chat panel?
     requestOpenWithChat: new Slot<AIChatParams>(),
+    requestSendWithChat: new Slot<AISendParams>(),
     requestInsertTemplate: new Slot<{
       template: string;
       mode: 'page' | 'edgeless';

packages/frontend/core/src/blocksuite/presets/ai/utils/extract.ts

@@ -0,0 +1,78 @@
+import type { EditorHost } from '@blocksuite/affine/block-std';
+import {
+  DocModeProvider,
+  isInsideEdgelessEditor,
+} from '@blocksuite/affine/blocks';
+
+import type { ChatContextValue } from '../chat-panel/chat-context';
+import {
+  getSelectedImagesAsBlobs,
+  getSelectedTextContent,
+  selectedToCanvas,
+} from './selection-utils';
+
+export async function extractContext(
+  host: EditorHost
+): Promise<Partial<ChatContextValue> | null> {
+  const docModeService = host.std.get(DocModeProvider);
+  const mode = docModeService.getEditorMode() || 'page';
+  if (mode === 'edgeless') {
+    return await extractOnEdgeless(host);
+  } else {
+    return await extractOnPage(host);
+  }
+}
+
+export async function extractOnEdgeless(
+  host: EditorHost
+): Promise<Partial<ChatContextValue> | null> {
+  if (!isInsideEdgelessEditor(host)) return null;
+
+  const canvas = await selectedToCanvas(host);
+  if (!canvas) return null;
+
+  const blob: Blob | null = await new Promise(resolve =>
+    canvas.toBlob(resolve)
+  );
+  if (!blob) return null;
+
+  return {
+    images: [new File([blob], 'selected.png')],
+  };
+}
+
+export async function extractOnPage(
+  host: EditorHost
+): Promise<Partial<ChatContextValue> | null> {
+  const text = await getSelectedTextContent(host, 'plain-text');
+  const images = await getSelectedImagesAsBlobs(host);
+  const hasText = text.length > 0;
+  const hasImages = images.length > 0;
+
+  if (hasText && !hasImages) {
+    const markdown = await getSelectedTextContent(host, 'markdown');
+    return {
+      quote: text,
+      markdown: markdown,
+    };
+  } else if (!hasText && hasImages && images.length === 1) {
+    host.command
+      .chain()
+      .tryAll(chain => [chain.getImageSelections()])
+      .getSelectedBlocks({
+        types: ['image'],
+      })
+      .run();
+    return {
+      images,
+    };
+  } else {
+    const markdown =
+      (await getSelectedTextContent(host, 'markdown')).trim() || '';
+    return {
+      quote: text,
+      markdown,
+      images,
+    };
+  }
+}

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

@@ -1,6 +1,8 @@
 import { TextRenderer } from './_common/components/text-renderer';
 import { AskAIButton } from './ai/_common/components/ask-ai-button';
+import { AskAIIcon } from './ai/_common/components/ask-ai-icon';
 import { AskAIPanel } from './ai/_common/components/ask-ai-panel';
+import { AskAIToolbarButton } from './ai/_common/components/ask-ai-toolbar';
 import { ChatActionList } from './ai/_common/components/chat-action-list';
 import { ChatCopyMore } from './ai/_common/components/copy-more';
 import { ChatPanel } from './ai/chat-panel';
@@ -36,7 +38,9 @@ import { ImagePlaceholder } from './blocks/ai-chat-block/components/image-placeh
 import { UserInfo } from './blocks/ai-chat-block/components/user-info';
 
 export function registerBlocksuitePresetsCustomComponents() {
+  customElements.define('ask-ai-icon', AskAIIcon);
   customElements.define('ask-ai-button', AskAIButton);
+  customElements.define('ask-ai-toolbar-button', AskAIToolbarButton);
   customElements.define('ask-ai-panel', AskAIPanel);
   customElements.define('chat-action-list', ChatActionList);
   customElements.define('chat-copy-more', ChatCopyMore);