refactor(core): use image preview component in chat (#10357)

[BS-2421](https://linear.app/affine-design/issue/BS-2421/chat-block-and-chat-panel-input-render-images-时存在内存泄露风险)

packages/frontend/core/src/blocksuite/presets/ai/_common/components/image-preview-grid.ts

@@ -0,0 +1,213 @@
+import { scrollbarStyle } from '@blocksuite/affine/blocks';
+import { CloseIcon } from '@blocksuite/icons/lit';
+import { css, html, LitElement } from 'lit';
+import { property, query, state } from 'lit/decorators.js';
+import { repeat } from 'lit/directives/repeat.js';
+
+export class ImagePreviewGrid extends LitElement {
+  static override styles = css`
+    .image-preview-wrapper {
+      overflow: hidden scroll;
+      max-height: 128px;
+    }
+
+    ${scrollbarStyle('.image-preview-wrapper')}
+
+    .images-container {
+      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;
+    }
+
+    .image-container 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);
+    }
+  `;
+
+  private readonly _urlMap = new Map<string, string>();
+  private readonly _urlRefCount = new Map<string, number>();
+
+  private _getFileKey(file: File) {
+    return `${file.name}-${file.size}-${file.lastModified}`;
+  }
+
+  private _disposeUrls() {
+    for (const [_, url] of this._urlMap.entries()) {
+      URL.revokeObjectURL(url);
+    }
+    this._urlRefCount.clear();
+    this._urlMap.clear();
+  }
+
+  /**
+   * get the object url of the file
+   * @param file - the file to get the url
+   * @returns the object url
+   */
+  private _getObjectUrl(file: File) {
+    const key = this._getFileKey(file);
+    let url = this._urlMap.get(key);
+
+    if (!url) {
+      // if the url is not in the map, create a new one
+      // and set the ref count to 0
+      url = URL.createObjectURL(file);
+      this._urlMap.set(key, url);
+      this._urlRefCount.set(url, 0);
+    }
+
+    // if the url is in the map, increment the ref count
+    const refCount = this._urlRefCount.get(url) || 0;
+    this._urlRefCount.set(url, refCount + 1);
+    return url;
+  }
+
+  /**
+   * decrement the reference count of the url
+   * when the reference count is 0, revoke the url
+   * @param url - the url to release
+   */
+  private readonly _releaseObjectUrl = (url: string) => {
+    const count = this._urlRefCount.get(url) || 0;
+    if (count <= 1) {
+      // when the last reference is released, revoke the url
+      URL.revokeObjectURL(url);
+      this._urlRefCount.delete(url);
+      // also delete the url from the map
+      for (const [key, value] of this._urlMap.entries()) {
+        if (value === url) {
+          this._urlMap.delete(key);
+          break;
+        }
+      }
+    } else {
+      // when the reference count is greater than 1, decrement the count
+      this._urlRefCount.set(url, count - 1);
+    }
+  };
+
+  private readonly _handleMouseEnter = (evt: MouseEvent, index: number) => {
+    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);
+    const top = Math.abs(parentRect.top - rect.top);
+    this.currentIndex = index;
+    if (!this.closeWrapper) return;
+    this.closeWrapper.style.display = 'flex';
+    this.closeWrapper.style.left = left + 'px';
+    this.closeWrapper.style.top = top + 'px';
+  };
+
+  private readonly _handleMouseLeave = () => {
+    if (!this.closeWrapper) return;
+    this.closeWrapper.style.display = 'none';
+    this.currentIndex = -1;
+  };
+
+  private readonly _handleDelete = () => {
+    if (this.currentIndex >= 0 && this.currentIndex < this.images.length) {
+      const file = this.images[this.currentIndex];
+      const url = this._getObjectUrl(file);
+      this._releaseObjectUrl(url);
+
+      this.onImageRemove?.(this.currentIndex);
+      this.currentIndex = -1;
+      if (!this.closeWrapper) return;
+      this.closeWrapper.style.display = 'none';
+    }
+  };
+
+  override disconnectedCallback() {
+    super.disconnectedCallback();
+    this._disposeUrls();
+  }
+
+  override render() {
+    return html`
+      <div class="image-preview-wrapper" @mouseleave=${this._handleMouseLeave}>
+        <div class="images-container">
+          ${repeat(
+            this.images,
+            image => this._getFileKey(image),
+            (image, index) => {
+              const url = this._getObjectUrl(image);
+              return html`
+                <div
+                  class="image-container"
+                  @error=${() => this._releaseObjectUrl(url)}
+                  @mouseenter=${(evt: MouseEvent) =>
+                    this._handleMouseEnter(evt, index)}
+                >
+                  <img src="${url}" alt="${image.name}" />
+                </div>
+              `;
+            }
+          )}
+        </div>
+        <div class="close-wrapper" @click=${this._handleDelete}>
+          ${CloseIcon()}
+        </div>
+      </div>
+    `;
+  }
+
+  @property({ type: Array })
+  accessor images: File[] = [];
+
+  @property({ attribute: false })
+  accessor onImageRemove: ((index: number) => void) | null = null;
+
+  @query('.close-wrapper')
+  accessor closeWrapper: HTMLDivElement | null = null;
+
+  @state()
+  accessor currentIndex = -1;
+}
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'image-preview-grid': ImagePreviewGrid;
+  }
+}

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

@@ -5,11 +5,7 @@ import {
   openFileOrFiles,
   unsafeCSSVarV2,
 } from '@blocksuite/affine/blocks';
-import {
+import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/utils';
-  assertExists,
-  SignalWatcher,
-  WithDisposable,
-} from '@blocksuite/affine/global/utils';
 import { ImageIcon, PublishIcon } from '@blocksuite/icons/lit';
 import { css, html, LitElement, nothing } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
@@ -169,38 +165,6 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
       }
     }
 
-    .chat-panel-images-wrapper {
-      overflow: hidden scroll;
-      max-height: 128px;
-
-      .chat-panel-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;
-          }
-        }
-      }
-    }
-
     .chat-panel-send svg rect {
       fill: var(--affine-primary-color);
     }
@@ -210,46 +174,17 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
     .chat-panel-send[aria-disabled='true'] svg rect {
       fill: var(--affine-text-disable-color);
     }
-
-    .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);
-    }
   `;
 
   @property({ attribute: false })
   accessor host!: EditorHost;
 
-  @query('.chat-panel-images')
+  @query('image-preview-grid')
-  accessor imagesWrapper!: HTMLDivElement;
+  accessor imagePreviewGrid: HTMLDivElement | null = null;
 
   @query('textarea')
   accessor textarea!: HTMLTextAreaElement;
 
-  @query('.close-wrapper')
-  accessor closeWrapper!: HTMLDivElement;
-
-  @state()
-  accessor curIndex = -1;
-
   @state()
   accessor isInputEmpty = true;
 
@@ -314,57 +249,11 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
     });
   }
 
-  private _renderImages(images: File[]) {
+  private readonly _handleImageRemove = (index: number) => {
-    return html`
+    const oldImages = this.chatContextValue.images;
-      <div
+    const newImages = oldImages.filter((_, i) => i !== index);
-        class="chat-panel-images-wrapper"
+    this.updateContext({ images: newImages });
-        @mouseleave=${() => {
+  };
-          this.closeWrapper.style.display = 'none';
-          this.curIndex = -1;
-        }}
-      >
-        <div class="chat-panel-images">
-          ${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();
-                  assertExists(ele.parentElement?.parentElement);
-                  const parentRect =
-                    ele.parentElement.parentElement.getBoundingClientRect();
-                  const left = Math.abs(rect.right - parentRect.left) - 8;
-                  const top = Math.abs(parentRect.top - rect.top);
-                  this.curIndex = index;
-                  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>
-        <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;
-              this.closeWrapper.style.display = 'none';
-            }
-          }}
-        >
-          ${CloseIcon}
-        </div>
-      </div>
-    `;
-  }
 
   private readonly _toggleNetworkSearch = (e: MouseEvent) => {
     e.preventDefault();
@@ -422,7 +311,14 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
           this.textarea.focus();
         }}
       >
-        ${hasImages ? this._renderImages(images) : nothing}
+        ${hasImages
+          ? html`
+              <image-preview-grid
+                .images=${images}
+                .onImageRemove=${this._handleImageRemove}
+              ></image-preview-grid>
+            `
+          : nothing}
         ${this.chatContextValue.quote
           ? html`<div class="chat-selection-quote">
               ${repeat(
@@ -448,7 +344,7 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
             this.isInputEmpty = !textarea.value.trim();
             textarea.style.height = 'auto';
             textarea.style.height = textarea.scrollHeight + 'px';
-            let imagesHeight = this.imagesWrapper?.scrollHeight ?? 0;
+            let imagesHeight = this.imagePreviewGrid?.scrollHeight ?? 0;
             if (imagesHeight) imagesHeight += 12;
             if (this.scrollHeight >= 200 + imagesHeight) {
               textarea.style.height = '148px';

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

@@ -9,15 +9,9 @@ import { ImageIcon, PublishIcon } from '@blocksuite/icons/lit';
 import { css, html, LitElement, nothing } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { classMap } from 'lit/directives/class-map.js';
-import { repeat } from 'lit/directives/repeat.js';
 
 import type { ChatMessage } from '../../../blocks';
-import {
+import { ChatAbortIcon, ChatClearIcon, ChatSendIcon } from '../_common/icons';
-  ChatAbortIcon,
-  ChatClearIcon,
-  ChatSendIcon,
-  CloseIcon,
-} from '../_common/icons';
 import type { AINetworkSearchConfig } from '../chat-panel/chat-config';
 import {
   PROMPT_NAME_AFFINE_AI,
@@ -78,30 +72,6 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {
         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;
-        }
-      }
-    }
 
     .chat-panel-send svg rect {
       fill: var(--affine-primary-color);
@@ -113,26 +83,6 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {
       fill: var(--affine-text-disable-color);
     }
 
-    .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;
@@ -198,7 +148,12 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {
         }
       </style>
       <div class="ai-chat-input">
-        ${hasImages ? this._renderImages(images) : nothing}
+        ${hasImages
+          ? html`<image-preview-grid
+              .images=${images}
+              .onImageRemove=${this._handleImageRemove}
+            ></image-preview-grid>`
+          : nothing}
         <textarea
           rows="1"
           placeholder="What are your thoughts?"
@@ -287,12 +242,6 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {
   @state()
   accessor _focused = false;
 
-  @query('.close-wrapper')
-  accessor closeWrapper: HTMLDivElement | null = null;
-
-  @state()
-  accessor _curIndex = -1;
-
   private _lastPromptName: string | null = null;
 
   private get _isNetworkActive() {
@@ -397,57 +346,11 @@ export class ChatBlockInput extends SignalWatcher(LitElement) {
     reportResponse('aborted:stop');
   };
 
-  private _renderImages(images: File[]) {
+  private readonly _handleImageRemove = (index: number) => {
-    return html`
+    const oldImages = this.chatContext.images;
-      <div
+    const newImages = oldImages.filter((_, i) => i !== index);
-        class="chat-input-images"
+    this.updateContext({ images: newImages });
-        @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 _onTextareaSend = async (e: MouseEvent | KeyboardEvent) => {
     e.preventDefault();

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

@@ -17,6 +17,7 @@ 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 { ImagePreviewGrid } from './ai/_common/components/image-preview-grid';
 import { ChatPanel } from './ai/chat-panel';
 import { ActionWrapper } from './ai/chat-panel/actions/action-wrapper';
 import { ChatText } from './ai/chat-panel/actions/chat-text';
@@ -50,6 +51,7 @@ export function registerBlocksuitePresetsCustomComponents() {
   customElements.define('ask-ai-panel', AskAIPanel);
   customElements.define('chat-action-list', ChatActionList);
   customElements.define('chat-copy-more', ChatCopyMore);
+  customElements.define('image-preview-grid', ImagePreviewGrid);
   customElements.define('action-wrapper', ActionWrapper);
   customElements.define('chat-text', ChatText);
   customElements.define('action-image-to-text', ActionImageToText);

tests/affine-cloud-copilot/e2e/copilot.spec.ts

@@ -806,7 +806,7 @@ test.describe('chat with block', () => {
           .waitForSelector('.ai-item-continue-with-ai')
           .then(i => i.click());
         await page
-          .waitForSelector('chat-panel-input .chat-panel-images')
+          .waitForSelector('chat-panel-input image-preview-grid')
           .then(el => el.waitForElementState('visible'));
       });
     });