feat(core): update ai add context button ui (#13172)

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

<img width="571" height="204" alt="截屏2025-07-11 17 33 01"
src="https://github.com/user-attachments/assets/3b7ed81f-1137-4c01-8fe2-9fe5ebf2adf3"
/>


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

* **New Features**
* Introduced a new component for adding context (images, documents,
tags, collections) to AI chat via a plus button and popover menu.
* Added notification feedback for duplicate chip additions and image
upload limits.
* Chips panel now supports collapsing and expanding for improved UI
control.

* **Improvements**
* Refactored chip management for better error handling, feedback, and
external control.
* Streamlined image and document uploads through a unified menu-driven
interface.
* Enhanced chip management methods with clearer naming and robust
synchronization.
* Updated chat input to delegate image upload and context additions to
the new add-context component.

* **Bug Fixes**
* Improved cancellation and cleanup of ongoing chip addition operations
to prevent conflicts.

* **Tests**
* Updated end-to-end tests to reflect the new menu-driven image upload
workflow and removed legacy checks.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Co-authored-by: fengmk2 <fengmk2@gmail.com>

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

@@ -20,10 +20,8 @@ import { property, state } from 'lit/decorators.js';
 import { keyed } from 'lit/directives/keyed.js';
 
 import { AffineIcon } from '../_common/icons';
-import type {
-  DocDisplayConfig,
-  SearchMenuConfig,
-} from '../components/ai-chat-chips';
+import type { SearchMenuConfig } from '../components/ai-chat-add-context';
+import type { DocDisplayConfig } from '../components/ai-chat-chips';
 import type { ChatContextValue } from '../components/ai-chat-content';
 import type {
   AINetworkSearchConfig,
@@ -385,6 +383,7 @@ export class ChatPanel extends SignalWatcher(
         .affineFeatureFlagService=${this.affineFeatureFlagService}
         .affineThemeService=${this.affineThemeService}
         .notificationService=${this.notificationService}
+        .affineWorkspaceDialogService=${this.affineWorkspaceDialogService}
       ></playground-content>
     `;
 

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

@@ -0,0 +1,93 @@
+import { createLitPortal } from '@blocksuite/affine/components/portal';
+import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
+import { ShadowlessElement } from '@blocksuite/affine/std';
+import { PlusIcon } from '@blocksuite/icons/lit';
+import { flip, offset } from '@floating-ui/dom';
+import { css, html } from 'lit';
+import { property, query } from 'lit/decorators.js';
+
+import type { ChatChip, DocDisplayConfig } from '../ai-chat-chips';
+import type { SearchMenuConfig } from './type';
+
+export class AIChatAddContext extends SignalWatcher(
+  WithDisposable(ShadowlessElement)
+) {
+  static override styles = css`
+    .ai-chat-add-context {
+      display: flex;
+      flex-shrink: 0;
+      flex-grow: 0;
+      align-items: center;
+      justify-content: center;
+      cursor: pointer;
+    }
+  `;
+
+  @property({ attribute: false })
+  accessor addChip!: (chip: ChatChip) => Promise<void>;
+
+  @property({ attribute: false })
+  accessor addImages!: (images: File[]) => void;
+
+  @property({ attribute: false })
+  accessor docDisplayConfig!: DocDisplayConfig;
+
+  @property({ attribute: false })
+  accessor searchMenuConfig!: SearchMenuConfig;
+
+  @property({ attribute: false })
+  accessor portalContainer: HTMLElement | null = null;
+
+  @query('.ai-chat-add-context')
+  accessor addButton!: HTMLDivElement;
+
+  private abortController: AbortController | null = null;
+
+  override render() {
+    return html`
+      <div
+        class="ai-chat-add-context"
+        data-testid="chat-panel-with-button"
+        @click=${this.toggleAddDocMenu}
+      >
+        ${PlusIcon()}
+      </div>
+    `;
+  }
+
+  private readonly toggleAddDocMenu = () => {
+    if (this.abortController) {
+      this.abortController.abort();
+      return;
+    }
+
+    this.abortController = new AbortController();
+    this.abortController.signal.addEventListener('abort', () => {
+      this.abortController = null;
+    });
+
+    createLitPortal({
+      template: html`
+        <chat-panel-add-popover
+          .addChip=${this.addChip}
+          .addImages=${this.addImages}
+          .searchMenuConfig=${this.searchMenuConfig}
+          .docDisplayConfig=${this.docDisplayConfig}
+          .abortController=${this.abortController}
+        ></chat-panel-add-popover>
+      `,
+      portalStyles: {
+        zIndex: 'var(--affine-z-index-popover)',
+      },
+      container: this.portalContainer ?? document.body,
+      computePosition: {
+        referenceElement: this.addButton,
+        placement: 'top-start',
+        middleware: [offset({ crossAxis: -30, mainAxis: 8 }), flip()],
+        autoUpdate: { animationFrame: true },
+      },
+      abortController: this.abortController,
+      closeOnClickAway: true,
+    });
+  };
+}

packages/frontend/core/src/blocksuite/ai/components/ai-chat-add-context/index.ts

@@ -0,0 +1,2 @@
+export * from './ai-chat-add-context';
+export * from './type';

packages/frontend/core/src/blocksuite/ai/components/ai-chat-add-context/type.ts

@@ -0,0 +1,24 @@
+import type {
+  SearchCollectionMenuAction,
+  SearchDocMenuAction,
+  SearchTagMenuAction,
+} from '@affine/core/modules/search-menu/services';
+import type { LinkedMenuGroup } from '@blocksuite/affine/widgets/linked-doc';
+
+export interface SearchMenuConfig {
+  getDocMenuGroup: (
+    query: string,
+    action: SearchDocMenuAction,
+    abortSignal: AbortSignal
+  ) => LinkedMenuGroup;
+  getTagMenuGroup: (
+    query: string,
+    action: SearchTagMenuAction,
+    abortSignal: AbortSignal
+  ) => LinkedMenuGroup;
+  getCollectionMenuGroup: (
+    query: string,
+    action: SearchCollectionMenuAction,
+    abortSignal: AbortSignal
+  ) => LinkedMenuGroup;
+}

packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/add-popover.ts

@@ -21,8 +21,8 @@ import { css, html, type TemplateResult } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
 
-import { MAX_IMAGE_COUNT } from '../ai-chat-input';
-import type { ChatChip, DocDisplayConfig, SearchMenuConfig } from './type';
+import type { SearchMenuConfig } from '../ai-chat-add-context';
+import type { ChatChip, DocDisplayConfig } from './type';
 
 enum AddPopoverMode {
   Default = 'default',
@@ -165,35 +165,31 @@ export class ChatPanelAddPopover extends SignalWatcher(
     const files = await openFilesWith();
     if (!files || files.length === 0) return;
 
+    this.abortController.abort();
     const images = files.filter(file => file.type.startsWith('image/'));
     if (images.length > 0) {
       this.addImages(images);
     }
 
     const others = files.filter(file => !file.type.startsWith('image/'));
-    for (const file of others) {
+    const addChipPromises = others.map(async file => {
       if (file.size > 50 * 1024 * 1024) {
         toast(`${file.name} is too large, please upload a file less than 50MB`);
-      } else {
-        await this.addChip({
-          file,
-          state: 'processing',
-        });
+        return;
       }
-    }
+      await this.addChip({
+        file,
+        state: 'processing',
+      });
+    });
+    await Promise.all(addChipPromises);
     this._track('file');
-    this.abortController.abort();
   };
 
   private readonly _addImageChip = async () => {
-    if (this.isImageUploadDisabled) return;
-
     const images = await openFilesWith('Images');
     if (!images) return;
-    if (this.uploadImageCount + images.length > MAX_IMAGE_COUNT) {
-      toast(`You can only upload up to ${MAX_IMAGE_COUNT} images`);
-      return;
-    }
+    this.abortController.abort();
     this.addImages(images);
   };
 
@@ -289,9 +285,6 @@ export class ChatPanelAddPopover extends SignalWatcher(
   @property({ attribute: 'data-testid', reflect: true })
   accessor testId: string = 'ai-search-input';
 
-  @property({ attribute: false })
-  accessor isImageUploadDisabled!: boolean;
-
   @property({ attribute: false })
   accessor uploadImageCount!: number;
 
@@ -498,31 +491,31 @@ export class ChatPanelAddPopover extends SignalWatcher(
   }
 
   private readonly _addDocChip = async (meta: DocMeta) => {
+    this.abortController.abort();
     await this.addChip({
       docId: meta.id,
       state: 'processing',
     });
     const mode = this.docDisplayConfig.getDocPrimaryMode(meta.id);
     this._track('doc', mode);
-    this.abortController.abort();
   };
 
   private readonly _addTagChip = async (tag: TagMeta) => {
+    this.abortController.abort();
     await this.addChip({
       tagId: tag.id,
       state: 'processing',
     });
     this._track('tags');
-    this.abortController.abort();
   };
 
   private readonly _addCollectionChip = async (collection: CollectionMeta) => {
+    this.abortController.abort();
     await this.addChip({
       collectionId: collection.id,
       state: 'processing',
     });
     this._track('collections');
-    this.abortController.abort();
   };
 
   private readonly _handleKeyDown = (event: KeyboardEvent) => {

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

@@ -3,7 +3,7 @@ import { createLitPortal } from '@blocksuite/affine/components/portal';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
 import { ShadowlessElement } from '@blocksuite/affine/std';
-import { MoreVerticalIcon, PlusIcon } from '@blocksuite/icons/lit';
+import { MoreVerticalIcon } from '@blocksuite/icons/lit';
 import { flip, offset } from '@floating-ui/dom';
 import { computed, type Signal, signal } from '@preact/signals-core';
 import { css, html, nothing, type PropertyValues } from 'lit';
@@ -11,16 +11,7 @@ import { property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
 import { isEqual } from 'lodash-es';
 
-import { AIProvider } from '../../provider';
-import type {
-  ChatChip,
-  CollectionChip,
-  DocChip,
-  DocDisplayConfig,
-  FileChip,
-  SearchMenuConfig,
-  TagChip,
-} from './type';
+import type { ChatChip, DocChip, DocDisplayConfig, FileChip } from './type';
 import {
   estimateTokenCount,
   getChipKey,
@@ -39,44 +30,46 @@ export class ChatPanelChips extends SignalWatcher(
   WithDisposable(ShadowlessElement)
 ) {
   static override styles = css`
-    .chips-wrapper {
+    .ai-chat-panel-chips {
       display: flex;
       gap: 8px;
       align-items: center;
       flex-wrap: wrap;
       padding: 4px 12px;
-    }
-    .add-button,
-    .collapse-button,
-    .more-candidate-button {
-      display: flex;
-      flex-shrink: 0;
-      flex-grow: 0;
-      align-items: center;
-      justify-content: center;
-      width: 24px;
-      height: 24px;
-      border: 0.5px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
-      border-radius: 4px;
-      box-sizing: border-box;
-      cursor: pointer;
-      font-size: 12px;
-      color: ${unsafeCSSVarV2('icon/primary')};
-    }
-    .add-button:hover,
-    .collapse-button:hover,
-    .more-candidate-button:hover {
-      background-color: ${unsafeCSSVarV2('layer/background/hoverOverlay')};
-    }
-    .more-candidate-button {
-      border-width: 1px;
-      border-style: dashed;
-      border-color: ${unsafeCSSVarV2('icon/tertiary')};
-      background: ${unsafeCSSVarV2('layer/background/secondary')};
-      color: ${unsafeCSSVarV2('icon/secondary')};
-    }
-    .more-candidate-button svg {
-      color: ${unsafeCSSVarV2('icon/secondary')};
+
+      .collapse-button,
+      .more-candidate-button {
+        display: flex;
+        flex-shrink: 0;
+        flex-grow: 0;
+        align-items: center;
+        justify-content: center;
+        width: 24px;
+        height: 24px;
+        border: 0.5px solid ${unsafeCSSVarV2('layer/insideBorder/border')};
+        border-radius: 4px;
+        box-sizing: border-box;
+        cursor: pointer;
+        font-size: 12px;
+        color: ${unsafeCSSVarV2('icon/primary')};
+      }
+
+      .collapse-button:hover,
+      .more-candidate-button:hover {
+        background-color: ${unsafeCSSVarV2('layer/background/hoverOverlay')};
+      }
+
+      .more-candidate-button {
+        border-width: 1px;
+        border-style: dashed;
+        border-color: ${unsafeCSSVarV2('icon/tertiary')};
+        background: ${unsafeCSSVarV2('layer/background/secondary')};
+        color: ${unsafeCSSVarV2('icon/secondary')};
+      }
+
+      .more-candidate-button svg {
+        color: ${unsafeCSSVarV2('icon/secondary')};
+      }
     }
   `;
 
@@ -86,38 +79,35 @@ export class ChatPanelChips extends SignalWatcher(
   accessor chips!: ChatChip[];
 
   @property({ attribute: false })
-  accessor createContextId!: () => Promise<string | undefined>;
+  accessor isCollapsed!: boolean;
 
   @property({ attribute: false })
-  accessor updateChips!: (chips: ChatChip[]) => void;
+  accessor addChip!: (chip: ChatChip) => Promise<void>;
 
   @property({ attribute: false })
-  accessor addImages!: (images: File[]) => void;
+  accessor updateChip!: (
+    chip: ChatChip,
+    options: Partial<DocChip | FileChip>
+  ) => void;
 
   @property({ attribute: false })
-  accessor pollContextDocsAndFiles!: () => void;
+  accessor removeChip!: (chip: ChatChip) => Promise<void>;
+
+  @property({ attribute: false })
+  accessor toggleCollapse!: () => void;
 
   @property({ attribute: false })
   accessor docDisplayConfig!: DocDisplayConfig;
 
-  @property({ attribute: false })
-  accessor searchMenuConfig!: SearchMenuConfig;
-
   @property({ attribute: false })
   accessor portalContainer: HTMLElement | null = null;
 
   @property({ attribute: 'data-testid', reflect: true })
   accessor testId = 'chat-panel-chips';
 
-  @query('.add-button')
-  accessor addButton!: HTMLDivElement;
-
   @query('.more-candidate-button')
   accessor moreCandidateButton!: HTMLDivElement;
 
-  @state()
-  accessor isCollapsed = false;
-
   @state()
   accessor referenceDocs: Signal<
     Array<{
@@ -144,14 +134,7 @@ export class ChatPanelChips extends SignalWatcher(
     const isCollapsed = this.isCollapsed && allChips.length > 1;
     const chips = isCollapsed ? allChips.slice(0, 1) : allChips;
 
-    return html`<div class="chips-wrapper">
-      <div
-        class="add-button"
-        data-testid="chat-panel-with-button"
-        @click=${this._toggleAddDocMenu}
-      >
-        ${PlusIcon()}
-      </div>
+    return html`<div class="ai-chat-panel-chips">
       ${repeat(
         chips,
         chip => getChipKey(chip),
@@ -159,9 +142,9 @@ export class ChatPanelChips extends SignalWatcher(
           if (isDocChip(chip)) {
             return html`<chat-panel-doc-chip
               .chip=${chip}
-              .addChip=${this._addChip}
-              .updateChip=${this._updateChip}
-              .removeChip=${this._removeChip}
+              .addChip=${this.addChip}
+              .updateChip=${this.updateChip}
+              .removeChip=${this.removeChip}
               .checkTokenLimit=${this._checkTokenLimit}
               .docDisplayConfig=${this.docDisplayConfig}
             ></chat-panel-doc-chip>`;
@@ -169,7 +152,7 @@ export class ChatPanelChips extends SignalWatcher(
           if (isFileChip(chip)) {
             return html`<chat-panel-file-chip
               .chip=${chip}
-              .removeChip=${this._removeChip}
+              .removeChip=${this.removeChip}
             ></chat-panel-file-chip>`;
           }
           if (isTagChip(chip)) {
@@ -180,7 +163,7 @@ export class ChatPanelChips extends SignalWatcher(
             return html`<chat-panel-tag-chip
               .chip=${chip}
               .tag=${tag}
-              .removeChip=${this._removeChip}
+              .removeChip=${this.removeChip}
             ></chat-panel-tag-chip>`;
           }
           if (isCollectionChip(chip)) {
@@ -193,7 +176,7 @@ export class ChatPanelChips extends SignalWatcher(
             return html`<chat-panel-collection-chip
               .chip=${chip}
               .collection=${collection}
-              .removeChip=${this._removeChip}
+              .removeChip=${this.removeChip}
             ></chat-panel-collection-chip>`;
           }
           return null;
@@ -208,7 +191,7 @@ export class ChatPanelChips extends SignalWatcher(
           </div>`
         : nothing}
       ${isCollapsed
-        ? html`<div class="collapse-button" @click=${this._toggleCollapse}>
+        ? html`<div class="collapse-button" @click=${this.toggleCollapse}>
             +${allChips.length - 1}
           </div>`
         : nothing}
@@ -227,14 +210,6 @@ export class ChatPanelChips extends SignalWatcher(
   }
 
   protected override updated(_changedProperties: PropertyValues): void {
-    if (
-      _changedProperties.has('chatContextValue') &&
-      _changedProperties.get('chatContextValue')?.status === 'loading' &&
-      this.isCollapsed === false
-    ) {
-      this.isCollapsed = true;
-    }
-
     if (_changedProperties.has('chips')) {
       this._updateReferenceDocs();
     }
@@ -245,46 +220,6 @@ export class ChatPanelChips extends SignalWatcher(
     this._cleanup?.();
   }
 
-  private readonly _toggleCollapse = () => {
-    this.isCollapsed = !this.isCollapsed;
-  };
-
-  private readonly _toggleAddDocMenu = () => {
-    if (this._abortController) {
-      this._abortController.abort();
-      return;
-    }
-
-    this._abortController = new AbortController();
-    this._abortController.signal.addEventListener('abort', () => {
-      this._abortController = null;
-    });
-
-    createLitPortal({
-      template: html`
-        <chat-panel-add-popover
-          .addChip=${this._addChip}
-          .addImages=${this.addImages}
-          .searchMenuConfig=${this.searchMenuConfig}
-          .docDisplayConfig=${this.docDisplayConfig}
-          .abortController=${this._abortController}
-        ></chat-panel-add-popover>
-      `,
-      portalStyles: {
-        zIndex: 'var(--affine-z-index-popover)',
-      },
-      container: this.portalContainer ?? document.body,
-      computePosition: {
-        referenceElement: this.addButton,
-        placement: 'top-start',
-        middleware: [offset({ crossAxis: -30, mainAxis: 8 }), flip()],
-        autoUpdate: { animationFrame: true },
-      },
-      abortController: this._abortController,
-      closeOnClickAway: true,
-    });
-  };
-
   private readonly _toggleMoreCandidatesMenu = () => {
     if (this._abortController) {
       this._abortController.abort();
@@ -303,7 +238,7 @@ export class ChatPanelChips extends SignalWatcher(
     createLitPortal({
       template: html`
         <chat-panel-candidates-popover
-          .addChip=${this._addChip}
+          .addChip=${this.addChip}
           .referenceDocs=${referenceDocs}
           .docDisplayConfig=${this.docDisplayConfig}
           .abortController=${this._abortController}
@@ -324,190 +259,6 @@ export class ChatPanelChips extends SignalWatcher(
     });
   };
 
-  private readonly _addChip = async (chip: ChatChip) => {
-    this.isCollapsed = false;
-    // remove the chip if it already exists
-    const chips = this._omitChip(this.chips, chip);
-    this.updateChips([...chips, chip]);
-    if (chips.length < this.chips.length) {
-      await this._removeFromContext(chip);
-    }
-    await this._addToContext(chip);
-    this.pollContextDocsAndFiles();
-  };
-
-  private readonly _updateChip = (
-    chip: ChatChip,
-    options: Partial<DocChip | FileChip>
-  ) => {
-    const index = this._findChipIndex(this.chips, chip);
-    if (index === -1) {
-      return;
-    }
-    const nextChip: ChatChip = {
-      ...chip,
-      ...options,
-    };
-    this.updateChips([
-      ...this.chips.slice(0, index),
-      nextChip,
-      ...this.chips.slice(index + 1),
-    ]);
-  };
-
-  private readonly _removeChip = async (chip: ChatChip) => {
-    const chips = this._omitChip(this.chips, chip);
-    this.updateChips(chips);
-    if (chips.length < this.chips.length) {
-      await this._removeFromContext(chip);
-    }
-  };
-
-  private readonly _addToContext = async (chip: ChatChip) => {
-    if (isDocChip(chip)) {
-      return await this._addDocToContext(chip);
-    }
-    if (isFileChip(chip)) {
-      return await this._addFileToContext(chip);
-    }
-    if (isTagChip(chip)) {
-      return await this._addTagToContext(chip);
-    }
-    if (isCollectionChip(chip)) {
-      return await this._addCollectionToContext(chip);
-    }
-    return null;
-  };
-
-  private readonly _addDocToContext = async (chip: DocChip) => {
-    try {
-      const contextId = await this.createContextId();
-      if (!contextId || !AIProvider.context) {
-        throw new Error('Context not found');
-      }
-      await AIProvider.context.addContextDoc({
-        contextId,
-        docId: chip.docId,
-      });
-    } catch (e) {
-      this._updateChip(chip, {
-        state: 'failed',
-        tooltip: e instanceof Error ? e.message : 'Add context doc error',
-      });
-    }
-  };
-
-  private readonly _addFileToContext = async (chip: FileChip) => {
-    try {
-      const contextId = await this.createContextId();
-      if (!contextId || !AIProvider.context) {
-        throw new Error('Context not found');
-      }
-      const contextFile = await AIProvider.context.addContextFile(chip.file, {
-        contextId,
-      });
-      this._updateChip(chip, {
-        state: contextFile.status,
-        blobId: contextFile.blobId,
-        fileId: contextFile.id,
-      });
-    } catch (e) {
-      this._updateChip(chip, {
-        state: 'failed',
-        tooltip: e instanceof Error ? e.message : 'Add context file error',
-      });
-    }
-  };
-
-  private readonly _addTagToContext = async (chip: TagChip) => {
-    try {
-      const contextId = await this.createContextId();
-      if (!contextId || !AIProvider.context) {
-        throw new Error('Context not found');
-      }
-      // TODO: server side docIds calculation
-      const docIds = this.docDisplayConfig.getTagPageIds(chip.tagId);
-      await AIProvider.context.addContextTag({
-        contextId,
-        tagId: chip.tagId,
-        docIds,
-      });
-      this._updateChip(chip, {
-        state: 'finished',
-      });
-    } catch (e) {
-      this._updateChip(chip, {
-        state: 'failed',
-        tooltip: e instanceof Error ? e.message : 'Add context tag error',
-      });
-    }
-  };
-
-  private readonly _addCollectionToContext = async (chip: CollectionChip) => {
-    try {
-      const contextId = await this.createContextId();
-      if (!contextId || !AIProvider.context) {
-        throw new Error('Context not found');
-      }
-      // TODO: server side docIds calculation
-      const docIds = this.docDisplayConfig.getCollectionPageIds(
-        chip.collectionId
-      );
-      await AIProvider.context.addContextCollection({
-        contextId,
-        collectionId: chip.collectionId,
-        docIds,
-      });
-      this._updateChip(chip, {
-        state: 'finished',
-      });
-    } catch (e) {
-      this._updateChip(chip, {
-        state: 'failed',
-        tooltip:
-          e instanceof Error ? e.message : 'Add context collection error',
-      });
-    }
-  };
-
-  private readonly _removeFromContext = async (
-    chip: ChatChip
-  ): Promise<boolean> => {
-    try {
-      const contextId = await this.createContextId();
-      if (!contextId || !AIProvider.context) {
-        return true;
-      }
-      if (isDocChip(chip)) {
-        return await AIProvider.context.removeContextDoc({
-          contextId,
-          docId: chip.docId,
-        });
-      }
-      if (isFileChip(chip) && chip.fileId) {
-        return await AIProvider.context.removeContextFile({
-          contextId,
-          fileId: chip.fileId,
-        });
-      }
-      if (isTagChip(chip)) {
-        return await AIProvider.context.removeContextTag({
-          contextId,
-          tagId: chip.tagId,
-        });
-      }
-      if (isCollectionChip(chip)) {
-        return await AIProvider.context.removeContextCollection({
-          contextId,
-          collectionId: chip.collectionId,
-        });
-      }
-      return true;
-    } catch {
-      return true;
-    }
-  };
-
   private readonly _checkTokenLimit = (
     newChip: DocChip,
     newTokenCount: number
@@ -544,44 +295,4 @@ export class ChatPanelChips extends SignalWatcher(
     this.referenceDocs = signal;
     this._cleanup = cleanup;
   };
-
-  private readonly _omitChip = (chips: ChatChip[], chip: ChatChip) => {
-    return chips.filter(item => {
-      if (isDocChip(chip)) {
-        return !isDocChip(item) || item.docId !== chip.docId;
-      }
-      if (isFileChip(chip)) {
-        return !isFileChip(item) || item.file !== chip.file;
-      }
-      if (isTagChip(chip)) {
-        return !isTagChip(item) || item.tagId !== chip.tagId;
-      }
-      if (isCollectionChip(chip)) {
-        return (
-          !isCollectionChip(item) || item.collectionId !== chip.collectionId
-        );
-      }
-      return true;
-    });
-  };
-
-  private readonly _findChipIndex = (chips: ChatChip[], chip: ChatChip) => {
-    return chips.findIndex(item => {
-      if (isDocChip(chip)) {
-        return isDocChip(item) && item.docId === chip.docId;
-      }
-      if (isFileChip(chip)) {
-        return isFileChip(item) && item.file === chip.file;
-      }
-      if (isTagChip(chip)) {
-        return isTagChip(item) && item.tagId === chip.tagId;
-      }
-      if (isCollectionChip(chip)) {
-        return (
-          isCollectionChip(item) && item.collectionId === chip.collectionId
-        );
-      }
-      return -1;
-    });
-  };
 }

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

@@ -1,11 +1,5 @@
 import type { TagMeta } from '@affine/core/components/page-list';
-import type {
-  SearchCollectionMenuAction,
-  SearchDocMenuAction,
-  SearchTagMenuAction,
-} from '@affine/core/modules/search-menu/services';
 import type { DocMeta, Store } from '@blocksuite/affine/store';
-import type { LinkedMenuGroup } from '@blocksuite/affine/widgets/linked-doc';
 import type { Signal } from '@preact/signals-core';
 
 export type ChipState = 'candidate' | 'processing' | 'finished' | 'failed';
@@ -75,21 +69,3 @@ export interface DocDisplayConfig {
   };
   getCollectionPageIds: (collectionId: string) => string[];
 }
-
-export interface SearchMenuConfig {
-  getDocMenuGroup: (
-    query: string,
-    action: SearchDocMenuAction,
-    abortSignal: AbortSignal
-  ) => LinkedMenuGroup;
-  getTagMenuGroup: (
-    query: string,
-    action: SearchTagMenuAction,
-    abortSignal: AbortSignal
-  ) => LinkedMenuGroup;
-  getCollectionMenuGroup: (
-    query: string,
-    action: SearchCollectionMenuAction,
-    abortSignal: AbortSignal
-  ) => LinkedMenuGroup;
-}

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

@@ -78,6 +78,42 @@ export function getChipKey(chip: ChatChip) {
   return null;
 }
 
+export function omitChip(chips: ChatChip[], chip: ChatChip) {
+  return chips.filter(item => {
+    if (isDocChip(chip)) {
+      return !isDocChip(item) || item.docId !== chip.docId;
+    }
+    if (isFileChip(chip)) {
+      return !isFileChip(item) || item.file !== chip.file;
+    }
+    if (isTagChip(chip)) {
+      return !isTagChip(item) || item.tagId !== chip.tagId;
+    }
+    if (isCollectionChip(chip)) {
+      return !isCollectionChip(item) || item.collectionId !== chip.collectionId;
+    }
+    return true;
+  });
+}
+
+export function findChipIndex(chips: ChatChip[], chip: ChatChip) {
+  return chips.findIndex(item => {
+    if (isDocChip(chip)) {
+      return isDocChip(item) && item.docId === chip.docId;
+    }
+    if (isFileChip(chip)) {
+      return isFileChip(item) && item.file === chip.file;
+    }
+    if (isTagChip(chip)) {
+      return isTagChip(item) && item.tagId === chip.tagId;
+    }
+    if (isCollectionChip(chip)) {
+      return isCollectionChip(item) && item.collectionId === chip.collectionId;
+    }
+    return -1;
+  });
+}
+
 export function estimateTokenCount(text: string): number {
   const chinese = text.match(/[\u4e00-\u9fa5]/g)?.length || 0;
   const english = text.replace(/[\u4e00-\u9fa5]/g, '');

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

@@ -12,20 +12,28 @@ import type {
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
-import { css, html } from 'lit';
+import type { NotificationService } from '@blocksuite/affine-shared/services';
+import { css, html, type PropertyValues } from 'lit';
 import { property, state } from 'lit/decorators.js';
 
 import { AIProvider } from '../../provider';
+import type { SearchMenuConfig } from '../ai-chat-add-context';
 import type {
   ChatChip,
   CollectionChip,
   DocChip,
   DocDisplayConfig,
   FileChip,
-  SearchMenuConfig,
   TagChip,
 } from '../ai-chat-chips';
-import { isCollectionChip, isDocChip, isTagChip } from '../ai-chat-chips';
+import {
+  findChipIndex,
+  isCollectionChip,
+  isDocChip,
+  isFileChip,
+  isTagChip,
+  omitChip,
+} from '../ai-chat-chips';
 import type {
   AIChatInputContext,
   AINetworkSearchConfig,
@@ -105,9 +113,15 @@ export class AIChatComposer extends SignalWatcher(
   @property({ attribute: false })
   accessor affineWorkspaceDialogService!: WorkspaceDialogService;
 
+  @property({ attribute: false })
+  accessor notificationService!: NotificationService;
+
   @state()
   accessor chips: ChatChip[] = [];
 
+  @state()
+  accessor isChipsCollapsed = false;
+
   @state()
   accessor embeddingCompleted = false;
 
@@ -121,11 +135,12 @@ export class AIChatComposer extends SignalWatcher(
     return html`
       <chat-panel-chips
         .chips=${this.chips}
-        .createContextId=${this._createContextId}
-        .updateChips=${this.updateChips}
-        .pollContextDocsAndFiles=${this._pollContextDocsAndFiles}
+        .isCollapsed=${this.isChipsCollapsed}
+        .addChip=${this.addChip}
+        .updateChip=${this.updateChip}
+        .removeChip=${this.removeChip}
+        .toggleCollapse=${this.toggleChipsCollapse}
         .docDisplayConfig=${this.docDisplayConfig}
-        .searchMenuConfig=${this.searchMenuConfig}
         .portalContainer=${this.portalContainer}
         .addImages=${this.addImages}
       ></chat-panel-chips>
@@ -136,15 +151,18 @@ export class AIChatComposer extends SignalWatcher(
         .docId=${this.docId}
         .session=${this.session}
         .chips=${this.chips}
+        .addChip=${this.addChip}
+        .addImages=${this.addImages}
         .createSession=${this.createSession}
         .chatContextValue=${this.chatContextValue}
         .updateContext=${this.updateContext}
         .networkSearchConfig=${this.networkSearchConfig}
         .reasoningConfig=${this.reasoningConfig}
         .docDisplayConfig=${this.docDisplayConfig}
+        .searchMenuConfig=${this.searchMenuConfig}
+        .portalContainer=${this.portalContainer}
         .onChatSuccess=${this.onChatSuccess}
         .trackOptions=${this.trackOptions}
-        .addImages=${this.addImages}
       ></ai-chat-input>
       <div class="chat-panel-footer">
         <ai-chat-composer-tip
@@ -165,7 +183,7 @@ export class AIChatComposer extends SignalWatcher(
 
   override connectedCallback() {
     super.connectedCallback();
-    this._initComposer().catch(console.error);
+    this.initComposer().catch(console.error);
   }
 
   override disconnectedCallback() {
@@ -174,6 +192,17 @@ export class AIChatComposer extends SignalWatcher(
     this._abortPollEmbeddingStatus();
   }
 
+  protected override willUpdate(changedProperties: PropertyValues): void {
+    if (
+      changedProperties.has('chatContextValue') &&
+      changedProperties.get('chatContextValue')?.status !== 'loading' &&
+      this.chatContextValue.status === 'loading' &&
+      this.isChipsCollapsed === false
+    ) {
+      this.isChipsCollapsed = true;
+    }
+  }
+
   private readonly _getContextId = async () => {
     if (this._contextId) {
       return this._contextId;
@@ -190,7 +219,7 @@ export class AIChatComposer extends SignalWatcher(
     return this._contextId;
   };
 
-  private readonly _createContextId = async () => {
+  private readonly createContextId = async () => {
     if (this._contextId) {
       return this._contextId;
     }
@@ -205,7 +234,7 @@ export class AIChatComposer extends SignalWatcher(
     return this._contextId;
   };
 
-  private readonly _initChips = async () => {
+  private readonly initChips = async () => {
     // context not initialized
     const sessionId = this.session?.sessionId;
     const contextId = await this._getContextId();
@@ -275,14 +304,206 @@ export class AIChatComposer extends SignalWatcher(
     this.chips = chips;
   };
 
+  private readonly updateChip = (
+    chip: ChatChip,
+    options: Partial<DocChip | FileChip>
+  ) => {
+    const index = findChipIndex(this.chips, chip);
+    if (index === -1) {
+      return;
+    }
+    const nextChip: ChatChip = {
+      ...chip,
+      ...options,
+    };
+    this.updateChips([
+      ...this.chips.slice(0, index),
+      nextChip,
+      ...this.chips.slice(index + 1),
+    ]);
+  };
+
+  private readonly addChip = async (chip: ChatChip) => {
+    this.isChipsCollapsed = false;
+    // if already exists
+    const index = findChipIndex(this.chips, chip);
+    if (index !== -1) {
+      this.notificationService.toast('chip already exists');
+      return;
+    }
+    this.updateChips([...this.chips, chip]);
+    await this.addToContext(chip);
+    await this.pollContextDocsAndFiles();
+  };
+
+  private readonly removeChip = async (chip: ChatChip) => {
+    const chips = omitChip(this.chips, chip);
+    this.updateChips(chips);
+    await this.removeFromContext(chip);
+  };
+
+  private readonly addToContext = async (chip: ChatChip) => {
+    if (isDocChip(chip)) {
+      return await this.addDocToContext(chip);
+    }
+    if (isFileChip(chip)) {
+      return await this.addFileToContext(chip);
+    }
+    if (isTagChip(chip)) {
+      return await this.addTagToContext(chip);
+    }
+    if (isCollectionChip(chip)) {
+      return await this.addCollectionToContext(chip);
+    }
+    return null;
+  };
+
+  private readonly addDocToContext = async (chip: DocChip) => {
+    try {
+      const contextId = await this.createContextId();
+      if (!contextId || !AIProvider.context) {
+        throw new Error('Context not found');
+      }
+      await AIProvider.context.addContextDoc({
+        contextId,
+        docId: chip.docId,
+      });
+    } catch (e) {
+      this.updateChip(chip, {
+        state: 'failed',
+        tooltip: e instanceof Error ? e.message : 'Add context doc error',
+      });
+    }
+  };
+
+  private readonly addFileToContext = async (chip: FileChip) => {
+    try {
+      const contextId = await this.createContextId();
+      if (!contextId || !AIProvider.context) {
+        throw new Error('Context not found');
+      }
+      const contextFile = await AIProvider.context.addContextFile(chip.file, {
+        contextId,
+      });
+      this.updateChip(chip, {
+        state: contextFile.status,
+        blobId: contextFile.blobId,
+        fileId: contextFile.id,
+      });
+    } catch (e) {
+      this.updateChip(chip, {
+        state: 'failed',
+        tooltip: e instanceof Error ? e.message : 'Add context file error',
+      });
+    }
+  };
+
+  private readonly addTagToContext = async (chip: TagChip) => {
+    try {
+      const contextId = await this.createContextId();
+      if (!contextId || !AIProvider.context) {
+        throw new Error('Context not found');
+      }
+      // TODO: server side docIds calculation
+      const docIds = this.docDisplayConfig.getTagPageIds(chip.tagId);
+      await AIProvider.context.addContextTag({
+        contextId,
+        tagId: chip.tagId,
+        docIds,
+      });
+      this.updateChip(chip, {
+        state: 'finished',
+      });
+    } catch (e) {
+      this.updateChip(chip, {
+        state: 'failed',
+        tooltip: e instanceof Error ? e.message : 'Add context tag error',
+      });
+    }
+  };
+
+  private readonly addCollectionToContext = async (chip: CollectionChip) => {
+    try {
+      const contextId = await this.createContextId();
+      if (!contextId || !AIProvider.context) {
+        throw new Error('Context not found');
+      }
+      // TODO: server side docIds calculation
+      const docIds = this.docDisplayConfig.getCollectionPageIds(
+        chip.collectionId
+      );
+      await AIProvider.context.addContextCollection({
+        contextId,
+        collectionId: chip.collectionId,
+        docIds,
+      });
+      this.updateChip(chip, {
+        state: 'finished',
+      });
+    } catch (e) {
+      this.updateChip(chip, {
+        state: 'failed',
+        tooltip:
+          e instanceof Error ? e.message : 'Add context collection error',
+      });
+    }
+  };
+
+  private readonly removeFromContext = async (
+    chip: ChatChip
+  ): Promise<boolean> => {
+    try {
+      const contextId = await this.createContextId();
+      if (!contextId || !AIProvider.context) {
+        return true;
+      }
+      if (isDocChip(chip)) {
+        return await AIProvider.context.removeContextDoc({
+          contextId,
+          docId: chip.docId,
+        });
+      }
+      if (isFileChip(chip) && chip.fileId) {
+        return await AIProvider.context.removeContextFile({
+          contextId,
+          fileId: chip.fileId,
+        });
+      }
+      if (isTagChip(chip)) {
+        return await AIProvider.context.removeContextTag({
+          contextId,
+          tagId: chip.tagId,
+        });
+      }
+      if (isCollectionChip(chip)) {
+        return await AIProvider.context.removeContextCollection({
+          contextId,
+          collectionId: chip.collectionId,
+        });
+      }
+      return true;
+    } catch {
+      return true;
+    }
+  };
+
+  private readonly toggleChipsCollapse = () => {
+    this.isChipsCollapsed = !this.isChipsCollapsed;
+  };
+
   private readonly addImages = (images: File[]) => {
     const oldImages = this.chatContextValue.images;
+    if (oldImages.length + images.length > MAX_IMAGE_COUNT) {
+      this.notificationService.toast(
+        `You can only upload up to ${MAX_IMAGE_COUNT} images`
+      );
+    }
     this.updateContext({
       images: [...oldImages, ...images].slice(0, MAX_IMAGE_COUNT),
     });
   };
 
-  private readonly _pollContextDocsAndFiles = async () => {
+  private readonly pollContextDocsAndFiles = async () => {
     const sessionId = this.session?.sessionId;
     const contextId = await this._getContextId();
     if (!sessionId || !contextId || !AIProvider.context) {
@@ -302,7 +523,7 @@ export class AIChatComposer extends SignalWatcher(
     );
   };
 
-  private readonly _pollEmbeddingStatus = async () => {
+  private readonly pollEmbeddingStatus = async () => {
     if (this._pollEmbeddingStatusAbortController) {
       this._pollEmbeddingStatusAbortController.abort();
     }
@@ -398,18 +619,18 @@ export class AIChatComposer extends SignalWatcher(
     this._pollEmbeddingStatusAbortController = null;
   };
 
-  private readonly _initComposer = async () => {
+  private readonly initComposer = async () => {
     const userId = (await AIProvider.userInfo)?.id;
     if (!userId || !this.session) return;
 
-    await this._initChips();
+    await this.initChips();
     const needPoll = this.chips.some(
       chip =>
         chip.state === 'processing' || isTagChip(chip) || isCollectionChip(chip)
     );
     if (needPoll) {
-      await this._pollContextDocsAndFiles();
+      await this.pollContextDocsAndFiles();
     }
-    await this._pollEmbeddingStatus();
+    await this.pollEmbeddingStatus();
   };
 }

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

@@ -25,7 +25,8 @@ import { styleMap } from 'lit/directives/style-map.js';
 import { HISTORY_IMAGE_ACTIONS } from '../../chat-panel/const';
 import { type AIChatParams, AIProvider } from '../../provider/ai-provider';
 import { extractSelectedContent } from '../../utils/extract';
-import type { DocDisplayConfig, SearchMenuConfig } from '../ai-chat-chips';
+import type { SearchMenuConfig } from '../ai-chat-add-context';
+import type { DocDisplayConfig } from '../ai-chat-chips';
 import type {
   AINetworkSearchConfig,
   AIReasoningConfig,
@@ -443,6 +444,7 @@ export class AIChatContent extends SignalWatcher(
         .docDisplayConfig=${this.docDisplayConfig}
         .searchMenuConfig=${this.searchMenuConfig}
         .affineWorkspaceDialogService=${this.affineWorkspaceDialogService}
+        .notificationService=${this.notificationService}
         .trackOptions=${{
           where: 'chat-panel',
           control: 'chat-send',

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

@@ -1,11 +1,9 @@
-import { toast } from '@affine/component';
 import type { CopilotChatHistoryFragment } from '@affine/graphql';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { unsafeCSSVar, unsafeCSSVarV2 } from '@blocksuite/affine/shared/theme';
-import { openFilesWith } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
-import { ArrowUpBigIcon, CloseIcon, ImageIcon } from '@blocksuite/icons/lit';
+import { ArrowUpBigIcon, CloseIcon } from '@blocksuite/icons/lit';
 import { css, html, nothing } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
@@ -16,6 +14,7 @@ import { type AIError, AIProvider, type AISendParams } from '../../provider';
 import { reportResponse } from '../../utils/action-reporter';
 import { readBlobAsURL } from '../../utils/image';
 import { mergeStreamObjects } from '../../utils/stream-objects';
+import type { SearchMenuConfig } from '../ai-chat-add-context';
 import type { ChatChip, DocDisplayConfig } from '../ai-chat-chips/type';
 import { isDocChip } from '../ai-chat-chips/utils';
 import {
@@ -23,7 +22,6 @@ import {
   isChatMessage,
   StreamObjectSchema,
 } from '../ai-chat-messages';
-import { MAX_IMAGE_COUNT } from './const';
 import type {
   AIChatInputContext,
   AINetworkSearchConfig,
@@ -335,6 +333,12 @@ export class AIChatInput extends SignalWatcher(
   @property({ attribute: false })
   accessor updateContext!: (context: Partial<AIChatInputContext>) => void;
 
+  @property({ attribute: false })
+  accessor addImages!: (images: File[]) => void;
+
+  @property({ attribute: false })
+  accessor addChip!: (chip: ChatChip) => Promise<void>;
+
   @property({ attribute: false })
   accessor networkSearchConfig!: AINetworkSearchConfig;
 
@@ -344,6 +348,9 @@ export class AIChatInput extends SignalWatcher(
   @property({ attribute: false })
   accessor docDisplayConfig!: DocDisplayConfig;
 
+  @property({ attribute: false })
+  accessor searchMenuConfig!: SearchMenuConfig;
+
   @property({ attribute: false })
   accessor isRootSession: boolean = true;
 
@@ -357,7 +364,7 @@ export class AIChatInput extends SignalWatcher(
   accessor testId = 'chat-panel-input-container';
 
   @property({ attribute: false })
-  accessor addImages!: (images: File[]) => void;
+  accessor portalContainer: HTMLElement | null = null;
 
   private get _isNetworkActive() {
     return (
@@ -370,10 +377,6 @@ export class AIChatInput extends SignalWatcher(
     return !!this.reasoningConfig.enabled.value;
   }
 
-  private get _isImageUploadDisabled() {
-    return this.chatContextValue.images.length >= MAX_IMAGE_COUNT;
-  }
-
   override connectedCallback() {
     super.connectedCallback();
     this._disposables.add(
@@ -453,14 +456,14 @@ export class AIChatInput extends SignalWatcher(
         data-testid="chat-panel-input"
       ></textarea>
       <div class="chat-panel-input-actions">
-        <div
-          class="chat-input-icon"
-          data-testid="chat-panel-input-image-upload"
-          aria-disabled=${this._isImageUploadDisabled}
-          @click=${this._uploadImageFiles}
-        >
-          ${ImageIcon()}
-          <affine-tooltip>Upload</affine-tooltip>
+        <div class="chat-input-icon">
+          <ai-chat-add-context
+            .addChip=${this.addChip}
+            .addImages=${this.addImages}
+            .docDisplayConfig=${this.docDisplayConfig}
+            .searchMenuConfig=${this.searchMenuConfig}
+            .portalContainer=${this.portalContainer}
+          ></ai-chat-add-context>
         </div>
         <div class="chat-input-footer-spacer"></div>
         <chat-input-preference
@@ -555,18 +558,6 @@ export class AIChatInput extends SignalWatcher(
     this.updateContext({ images: newImages });
   };
 
-  private readonly _uploadImageFiles = async (_e: MouseEvent) => {
-    if (this._isImageUploadDisabled) return;
-
-    const images = await openFilesWith('Images');
-    if (!images) return;
-    if (this.chatContextValue.images.length + images.length > MAX_IMAGE_COUNT) {
-      toast(`You can only upload up to ${MAX_IMAGE_COUNT} images`);
-      return;
-    }
-    this.addImages(images);
-  };
-
   private readonly _onTextareaSend = async (e: MouseEvent | KeyboardEvent) => {
     e.preventDefault();
     e.stopPropagation();

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

@@ -1,3 +1,4 @@
+import type { WorkspaceDialogService } from '@affine/core/modules/dialogs';
 import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import type { AppThemeService } from '@affine/core/modules/theme';
 import type {
@@ -19,7 +20,8 @@ import { throttle } from 'lodash-es';
 import type { AppSidebarConfig } from '../../chat-panel/chat-config';
 import { HISTORY_IMAGE_ACTIONS } from '../../chat-panel/const';
 import { AIProvider } from '../../provider';
-import type { DocDisplayConfig, SearchMenuConfig } from '../ai-chat-chips';
+import type { SearchMenuConfig } from '../ai-chat-add-context';
+import type { DocDisplayConfig } from '../ai-chat-chips';
 import type { ChatContextValue } from '../ai-chat-content';
 import type {
   AINetworkSearchConfig,
@@ -165,6 +167,9 @@ export class PlaygroundChat extends SignalWatcher(
   @property({ attribute: false })
   accessor affineThemeService!: AppThemeService;
 
+  @property({ attribute: false })
+  accessor affineWorkspaceDialogService!: WorkspaceDialogService;
+
   @property({ attribute: false })
   accessor notificationService!: NotificationService;
 
@@ -351,6 +356,8 @@ export class PlaygroundChat extends SignalWatcher(
         .playgroundConfig=${this.playgroundConfig}
         .docDisplayConfig=${this.docDisplayConfig}
         .searchMenuConfig=${this.searchMenuConfig}
+        .notificationService=${this.notificationService}
+        .affineWorkspaceDialogService=${this.affineWorkspaceDialogService}
       ></ai-chat-composer>
     </div>`;
   }

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

@@ -12,7 +12,8 @@ import { repeat } from 'lit/directives/repeat.js';
 
 import type { AppSidebarConfig } from '../../chat-panel/chat-config';
 import { AIProvider } from '../../provider';
-import type { DocDisplayConfig, SearchMenuConfig } from '../ai-chat-chips';
+import type { SearchMenuConfig } from '../ai-chat-add-context';
+import type { DocDisplayConfig } from '../ai-chat-chips';
 import type {
   AINetworkSearchConfig,
   AIPlaygroundConfig,

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

@@ -26,6 +26,7 @@ import { ChatMessageAction } from './chat-panel/message/action';
 import { ChatMessageAssistant } from './chat-panel/message/assistant';
 import { ChatMessageUser } from './chat-panel/message/user';
 import { ChatPanelSplitView } from './chat-panel/split-view';
+import { AIChatAddContext } from './components/ai-chat-add-context';
 import { ChatPanelAddPopover } from './components/ai-chat-chips/add-popover';
 import { ChatPanelCandidatesPopover } from './components/ai-chat-chips/candidates-popover';
 import { ChatPanelChips } from './components/ai-chat-chips/chat-panel-chips';
@@ -138,6 +139,7 @@ export function registerAIEffects() {
   customElements.define('ai-chat-messages', AIChatMessages);
   customElements.define('chat-panel', ChatPanel);
   customElements.define('ai-chat-input', AIChatInput);
+  customElements.define('ai-chat-add-context', AIChatAddContext);
   customElements.define(
     'ai-chat-embedding-status-tooltip',
     AIChatEmbeddingStatusTooltip

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

@@ -29,10 +29,8 @@ import {
   queryHistoryMessages,
 } from '../_common/chat-actions-handle';
 import { type AIChatBlockModel } from '../blocks';
-import type {
-  DocDisplayConfig,
-  SearchMenuConfig,
-} from '../components/ai-chat-chips';
+import type { SearchMenuConfig } from '../components/ai-chat-add-context';
+import type { DocDisplayConfig } from '../components/ai-chat-chips';
 import type {
   AINetworkSearchConfig,
   AIReasoningConfig,
@@ -609,6 +607,7 @@ export class AIChatBlockPeekView extends LitElement {
         .docDisplayConfig=${this.docDisplayConfig}
         .searchMenuConfig=${this.searchMenuConfig}
         .affineWorkspaceDialogService=${this.affineWorkspaceDialogService}
+        .notificationService=${notificationService}
         .onChatSuccess=${this._onChatSuccess}
         .trackOptions=${{
           where: 'ai-chat-block',

tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts

@@ -260,8 +260,12 @@ export class ChatPanelUtils {
     });
 
     const fileChooserPromise = page.waitForEvent('filechooser');
-    // Open file upload dialog
-    await page.getByTestId('chat-panel-input-image-upload').click();
+    const withButton = page.getByTestId('chat-panel-with-button');
+    await withButton.hover();
+    await withButton.click({ delay: 200 });
+    const withMenu = page.getByTestId('ai-add-popover');
+    await withMenu.waitFor({ state: 'visible' });
+    await withMenu.getByTestId('ai-chat-with-images').click();
 
     const fileChooser = await fileChooserPromise;
     await fileChooser.setFiles(images);
@@ -369,10 +373,4 @@ export class ChatPanelUtils {
     const networkSearch = await page.getByTestId('chat-network-search');
     return (await networkSearch.getAttribute('aria-disabled')) === 'false';
   }
-
-  public static async isImageUploadEnabled(page: Page) {
-    const imageUpload = await page.getByTestId('chat-panel-input-image-upload');
-    const disabled = await imageUpload.getAttribute('data-disabled');
-    return disabled === 'false';
-  }
 }