feat(editor): comment for edgeless element (#13098)

#### PR Dependency Tree


* **PR #13098** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

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

* **New Features**
* Added support for comments on graphical elements, allowing users to
comment on both blocks and graphical elements within surfaces.
* Enhanced comment previews to include graphical elements in selection
summaries.
* Improved editor navigation to focus on commented graphical elements in
addition to blocks and inline texts.

* **Bug Fixes**
* Updated comment highlighting and management to consistently use the
new comment manager across all block and element types.

* **Refactor**
* Renamed and extended the comment manager to handle both block and
element comments.
* Streamlined toolbar configurations by removing outdated comment button
entries and adding a consolidated comment button in the root toolbar.

* **Tests**
* Disabled the mock comment provider integration in the test editor
environment to refine testing setup.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

blocksuite/affine/blocks/attachment/src/attachment-block.ts

@@ -17,7 +17,7 @@ import {
   AttachmentBlockStyles,
 } from '@blocksuite/affine-model';
 import {
-  BlockCommentManager,
+  BlockElementCommentManager,
   CitationProvider,
   DocModeProvider,
   FileSizeLimitProvider,
@@ -96,7 +96,7 @@ export class AttachmentBlockComponent extends CaptionedBlockComponent<Attachment
   get isCommentHighlighted() {
     return (
       this.std
-        .getOptional(BlockCommentManager)
+        .getOptional(BlockElementCommentManager)
         ?.isBlockCommentHighlighted(this.model) ?? false
     );
   }

blocksuite/affine/blocks/attachment/src/configs/toolbar.ts

@@ -10,7 +10,6 @@ import {
 } from '@blocksuite/affine-shared/consts';
 import {
   ActionPlacement,
-  blockCommentToolbarButton,
   type ToolbarAction,
   type ToolbarActionGroup,
   type ToolbarModuleConfig,
@@ -241,10 +240,6 @@ const builtinToolbarConfig = {
     replaceAction,
     downloadAction,
     captionAction,
-    {
-      id: 'f.comment',
-      ...blockCommentToolbarButton,
-    },
     {
       placement: ActionPlacement.More,
       id: 'a.clipboard',

blocksuite/affine/blocks/bookmark/src/bookmark-block.ts

@@ -8,7 +8,7 @@ import type {
 } from '@blocksuite/affine-model';
 import { ImageProxyService } from '@blocksuite/affine-shared/adapters';
 import {
-  BlockCommentManager,
+  BlockElementCommentManager,
   CitationProvider,
   DocModeProvider,
   LinkPreviewServiceIdentifier,
@@ -132,7 +132,7 @@ export class BookmarkBlockComponent extends CaptionedBlockComponent<BookmarkBloc
   get isCommentHighlighted() {
     return (
       this.std
-        .getOptional(BlockCommentManager)
+        .getOptional(BlockElementCommentManager)
         ?.isBlockCommentHighlighted(this.model) ?? false
     );
   }

blocksuite/affine/blocks/bookmark/src/configs/toolbar.ts

@@ -17,7 +17,6 @@ import {
 } from '@blocksuite/affine-shared/consts';
 import {
   ActionPlacement,
-  blockCommentToolbarButton,
   EmbedIframeService,
   EmbedOptionProvider,
   type LinkEventType,
@@ -289,10 +288,6 @@ const builtinToolbarConfig = {
       },
     } satisfies ToolbarActionGroup<ToolbarAction>,
     captionAction,
-    {
-      id: 'e.comment',
-      ...blockCommentToolbarButton,
-    },
     {
       placement: ActionPlacement.More,
       id: 'a.clipboard',

blocksuite/affine/blocks/code/src/code-block.ts

@@ -6,7 +6,7 @@ import {
   EDGELESS_TOP_CONTENTEDITABLE_SELECTOR,
 } from '@blocksuite/affine-shared/consts';
 import {
-  BlockCommentManager,
+  BlockElementCommentManager,
   DocModeProvider,
   NotificationProvider,
 } from '@blocksuite/affine-shared/services';
@@ -394,7 +394,7 @@ export class CodeBlockComponent extends CaptionedBlockComponent<CodeBlockModel>
   get isCommentHighlighted() {
     return (
       this.std
-        .getOptional(BlockCommentManager)
+        .getOptional(BlockElementCommentManager)
         ?.isBlockCommentHighlighted(this.model) ?? false
     );
   }

blocksuite/affine/blocks/database/src/database-block.ts

@@ -10,7 +10,7 @@ import { toast } from '@blocksuite/affine-components/toast';
 import type { DatabaseBlockModel } from '@blocksuite/affine-model';
 import { EDGELESS_TOP_CONTENTEDITABLE_SELECTOR } from '@blocksuite/affine-shared/consts';
 import {
-  BlockCommentManager,
+  BlockElementCommentManager,
   CommentProviderIdentifier,
   DocModeProvider,
   NotificationProvider,
@@ -316,7 +316,7 @@ export class DatabaseBlockComponent extends CaptionedBlockComponent<DatabaseBloc
   get isCommentHighlighted() {
     return (
       this.std
-        .getOptional(BlockCommentManager)
+        .getOptional(BlockElementCommentManager)
         ?.isBlockCommentHighlighted(this.model) ?? false
     );
   }

blocksuite/affine/blocks/embed-doc/src/embed-linked-doc-block/configs/toolbar.ts

@@ -11,7 +11,6 @@ import {
 } from '@blocksuite/affine-shared/consts';
 import {
   ActionPlacement,
-  blockCommentToolbarButton,
   DocDisplayMetaProvider,
   EditorSettingProvider,
   type LinkEventType,
@@ -306,10 +305,6 @@ const builtinToolbarConfig = {
       },
     } satisfies ToolbarActionGroup<ToolbarAction>,
     captionAction,
-    {
-      id: 'e.comment',
-      ...blockCommentToolbarButton,
-    },
     {
       placement: ActionPlacement.More,
       id: 'a.clipboard',

blocksuite/affine/blocks/embed-doc/src/embed-synced-doc-block/configs/toolbar.ts

@@ -16,7 +16,6 @@ import {
 import { REFERENCE_NODE } from '@blocksuite/affine-shared/consts';
 import {
   ActionPlacement,
-  blockCommentToolbarButton,
   EditorSettingProvider,
   type LinkEventType,
   type OpenDocMode,
@@ -226,10 +225,6 @@ const builtinToolbarConfig = {
     openDocActionGroup,
     conversionsActionGroup,
     captionAction,
-    {
-      id: 'e.comment',
-      ...blockCommentToolbarButton,
-    },
     {
       placement: ActionPlacement.More,
       id: 'a.clipboard',

blocksuite/affine/blocks/embed/src/common/embed-block-element.ts

@@ -9,7 +9,7 @@ import {
   EMBED_CARD_WIDTH,
 } from '@blocksuite/affine-shared/consts';
 import {
-  BlockCommentManager,
+  BlockElementCommentManager,
   DocModeProvider,
 } from '@blocksuite/affine-shared/services';
 import { unsafeCSSVarV2 } from '@blocksuite/affine-shared/theme';
@@ -65,7 +65,7 @@ export class EmbedBlockComponent<
   get isCommentHighlighted() {
     return (
       this.std
-        .getOptional(BlockCommentManager)
+        .getOptional(BlockElementCommentManager)
         ?.isBlockCommentHighlighted(this.model) ?? false
     );
   }

blocksuite/affine/blocks/embed/src/configs/toolbar.ts

@@ -13,7 +13,6 @@ import {
 } from '@blocksuite/affine-shared/consts';
 import {
   ActionPlacement,
-  blockCommentToolbarButton,
   EmbedOptionProvider,
   type LinkEventType,
   type ToolbarAction,
@@ -349,10 +348,6 @@ function createBuiltinToolbarConfigForExternal(
           });
         },
       },
-      {
-        id: 'e.comment',
-        ...blockCommentToolbarButton,
-      },
       {
         placement: ActionPlacement.More,
         id: 'a.clipboard',

blocksuite/affine/blocks/image/src/configs/toolbar.ts

@@ -1,7 +1,6 @@
 import { ImageBlockModel } from '@blocksuite/affine-model';
 import {
   ActionPlacement,
-  blockCommentToolbarButton,
   type ToolbarModuleConfig,
   ToolbarModuleExtension,
 } from '@blocksuite/affine-shared/services';
@@ -50,10 +49,6 @@ const builtinToolbarConfig = {
         });
       },
     },
-    {
-      id: 'c.comment',
-      ...blockCommentToolbarButton,
-    },
     {
       placement: ActionPlacement.More,
       id: 'a.clipboard',
@@ -146,10 +141,6 @@ const builtinSurfaceToolbarConfig = {
         });
       },
     },
-    {
-      id: 'c.comment',
-      ...blockCommentToolbarButton,
-    },
   ],
 
   when: ctx => ctx.getSurfaceModelsByType(ImageBlockModel).length === 1,

blocksuite/affine/blocks/image/src/image-block.ts

@@ -6,7 +6,7 @@ import { ResourceController } from '@blocksuite/affine-components/resource';
 import type { ImageBlockModel } from '@blocksuite/affine-model';
 import { ImageSelection } from '@blocksuite/affine-shared/selection';
 import {
-  BlockCommentManager,
+  BlockElementCommentManager,
   ToolbarRegistryIdentifier,
 } from '@blocksuite/affine-shared/services';
 import { formatSize } from '@blocksuite/affine-shared/utils';
@@ -71,7 +71,7 @@ export class ImageBlockComponent extends CaptionedBlockComponent<ImageBlockModel
   get isCommentHighlighted() {
     return (
       this.std
-        .getOptional(BlockCommentManager)
+        .getOptional(BlockElementCommentManager)
         ?.isBlockCommentHighlighted(this.model) ?? false
     );
   }

blocksuite/affine/blocks/paragraph/src/paragraph-block.ts

@@ -8,7 +8,7 @@ import {
   EDGELESS_TOP_CONTENTEDITABLE_SELECTOR,
 } from '@blocksuite/affine-shared/consts';
 import {
-  BlockCommentManager,
+  BlockElementCommentManager,
   CitationProvider,
   DocModeProvider,
 } from '@blocksuite/affine-shared/services';
@@ -112,7 +112,7 @@ export class ParagraphBlockComponent extends CaptionedBlockComponent<ParagraphBl
   get isCommentHighlighted() {
     return (
       this.std
-        .getOptional(BlockCommentManager)
+        .getOptional(BlockElementCommentManager)
         ?.isBlockCommentHighlighted(this.model) ?? false
     );
   }

blocksuite/affine/blocks/root/src/edgeless/configs/toolbar/misc.ts

@@ -17,6 +17,7 @@ import {
 } from '@blocksuite/affine-model';
 import {
   ActionPlacement,
+  blockCommentToolbarButton,
   type ElementLockEvent,
   type ToolbarAction,
   type ToolbarContext,
@@ -305,6 +306,12 @@ export const builtinMiscToolbarConfig = {
       },
     },
 
+    {
+      placement: ActionPlacement.End,
+      id: 'c.comment',
+      ...blockCommentToolbarButton,
+    },
+
     // More actions
     ...moreActions.map(action => ({
       ...action,

blocksuite/affine/blocks/surface-ref/src/configs/toolbar.ts

@@ -5,7 +5,6 @@ import {
 } from '@blocksuite/affine-shared/commands';
 import {
   ActionPlacement,
-  blockCommentToolbarButton,
   type ToolbarModuleConfig,
 } from '@blocksuite/affine-shared/services';
 import { CaptionIcon, CopyIcon, DeleteIcon } from '@blocksuite/icons/lit';
@@ -62,10 +61,6 @@ export const surfaceRefToolbarModuleConfig: ToolbarModuleConfig = {
         surfaceRefBlock.captionElement.show();
       },
     },
-    {
-      id: 'e.comment',
-      ...blockCommentToolbarButton,
-    },
     {
       id: 'a.clipboard',
       placement: ActionPlacement.More,

blocksuite/affine/blocks/surface-ref/src/surface-ref-block.ts

@@ -13,7 +13,7 @@ import {
   type SurfaceRefBlockModel,
 } from '@blocksuite/affine-model';
 import {
-  BlockCommentManager,
+  BlockElementCommentManager,
   DocModeProvider,
   EditPropsStore,
   type OpenDocMode,
@@ -145,7 +145,7 @@ export class SurfaceRefBlockComponent extends BlockComponent<SurfaceRefBlockMode
   get isCommentHighlighted() {
     return (
       this.std
-        .getOptional(BlockCommentManager)
+        .getOptional(BlockElementCommentManager)
         ?.isBlockCommentHighlighted(this.model) ?? false
     );
   }

blocksuite/affine/foundation/src/view.ts

@@ -9,7 +9,7 @@ import {
 } from '@blocksuite/affine-ext-loader';
 import {
   AutoClearSelectionService,
-  BlockCommentManager,
+  BlockElementCommentManager,
   CitationService,
   DefaultOpenDocExtension,
   DNDAPIExtension,
@@ -79,7 +79,7 @@ export class FoundationViewExtension extends ViewExtensionProvider<FoundationVie
       LinkPreviewCache,
       LinkPreviewService,
       CitationService,
-      BlockCommentManager,
+      BlockElementCommentManager,
     ]);
     context.register(clipboardConfigs);
     if (this.isEdgeless(context.scope)) {

blocksuite/affine/gfx/text/src/toolbar/actions.ts

@@ -17,9 +17,9 @@ import {
   TextAlign,
   type TextStyleProps,
 } from '@blocksuite/affine-model';
-import type {
+import {
-  ToolbarActions,
+  type ToolbarActions,
-  ToolbarContext,
+  type ToolbarContext,
 } from '@blocksuite/affine-shared/services';
 import {
   getMostCommonResolvedValue,

blocksuite/affine/inlines/comment/src/inline-comment-manager.ts

@@ -1,7 +1,7 @@
 import { getInlineEditorByModel } from '@blocksuite/affine-rich-text';
 import { getSelectedBlocksCommand } from '@blocksuite/affine-shared/commands';
 import {
-  BlockCommentManager,
+  BlockElementCommentManager,
   type CommentId,
   CommentProviderIdentifier,
   findAllCommentedBlocks,
@@ -78,7 +78,7 @@ export class InlineCommentManager extends LifeCycleWatcher {
     // which means the comment may be removed or resolved in provider side
     difference(commentsInEditor, commentsInProvider).forEach(comment => {
       this._handleDeleteAndResolve(comment);
-      this.std.get(BlockCommentManager).handleDeleteAndResolve(comment);
+      this.std.get(BlockElementCommentManager).handleDeleteAndResolve(comment);
     });
   }
 

blocksuite/affine/shared/src/services/comment-service/block-comment-manager.ts

@@ -1,118 +0,0 @@
-import { DividerBlockModel } from '@blocksuite/affine-model';
-import { DisposableGroup } from '@blocksuite/global/disposable';
-import {
-  BlockSelection,
-  LifeCycleWatcher,
-  TextSelection,
-} from '@blocksuite/std';
-import type { BaseSelection, BlockModel } from '@blocksuite/store';
-import { signal } from '@preact/signals-core';
-
-import { getSelectedBlocksCommand } from '../../commands';
-import { ImageSelection } from '../../selection';
-import { matchModels } from '../../utils';
-import { type CommentId, CommentProviderIdentifier } from './comment-provider';
-import { findCommentedBlocks } from './utils';
-
-export class BlockCommentManager extends LifeCycleWatcher {
-  static override key = 'block-comment-manager';
-
-  private readonly _highlightedCommentId$ = signal<CommentId | null>(null);
-
-  private readonly _disposables = new DisposableGroup();
-
-  private get _provider() {
-    return this.std.getOptional(CommentProviderIdentifier);
-  }
-
-  isBlockCommentHighlighted(
-    block: BlockModel<{ comments?: Record<CommentId, boolean> }>
-  ) {
-    const comments = block.props.comments;
-    if (!comments) return false;
-    return (
-      this._highlightedCommentId$.value !== null &&
-      Object.keys(comments).includes(this._highlightedCommentId$.value)
-    );
-  }
-
-  override mounted() {
-    const provider = this._provider;
-    if (!provider) return;
-
-    this._disposables.add(provider.onCommentAdded(this._handleAddComment));
-    this._disposables.add(
-      provider.onCommentDeleted(this.handleDeleteAndResolve)
-    );
-    this._disposables.add(
-      provider.onCommentResolved(this.handleDeleteAndResolve)
-    );
-    this._disposables.add(
-      provider.onCommentHighlighted(this._handleHighlightComment)
-    );
-  }
-
-  override unmounted() {
-    this._disposables.dispose();
-  }
-
-  private readonly _handleAddComment = (
-    id: CommentId,
-    selections: BaseSelection[]
-  ) => {
-    const blocksFromTextRange = selections
-      .filter((s): s is TextSelection => s.is(TextSelection))
-      .map(s => {
-        const [_, { selectedBlocks }] = this.std.command.exec(
-          getSelectedBlocksCommand,
-          {
-            textSelection: s,
-          }
-        );
-        if (!selectedBlocks) return [];
-        return selectedBlocks.map(b => b.model).filter(m => !m.text);
-      });
-
-    const needCommentBlocks = [
-      ...blocksFromTextRange.flat(),
-      ...selections
-        .filter(s => s instanceof BlockSelection || s instanceof ImageSelection)
-        .map(({ blockId }) => this.std.store.getModelById(blockId))
-        .filter(
-          (m): m is BlockModel =>
-            m !== null && !matchModels(m, [DividerBlockModel])
-        ),
-    ];
-
-    if (needCommentBlocks.length === 0) return;
-
-    this.std.store.withoutTransact(() => {
-      needCommentBlocks.forEach(block => {
-        const comments = (
-          'comments' in block.props &&
-          typeof block.props.comments === 'object' &&
-          block.props.comments !== null
-            ? block.props.comments
-            : {}
-        ) as Record<CommentId, boolean>;
-
-        this.std.store.updateBlock(block, {
-          comments: { [id]: true, ...comments },
-        });
-      });
-    });
-  };
-
-  readonly handleDeleteAndResolve = (id: CommentId) => {
-    const commentedBlocks = findCommentedBlocks(this.std.store, id);
-    this.std.store.withoutTransact(() => {
-      commentedBlocks.forEach(block => {
-        delete block.props.comments[id];
-      });
-    });
-  };
-
-  private readonly _handleHighlightComment = (id: CommentId | null) => {
-    this._highlightedCommentId$.value = id;
-  };
-}

blocksuite/affine/shared/src/services/comment-service/block-element-comment-manager.ts

@@ -0,0 +1,169 @@
+import { DividerBlockModel } from '@blocksuite/affine-model';
+import { DisposableGroup } from '@blocksuite/global/disposable';
+import {
+  BlockSelection,
+  LifeCycleWatcher,
+  SurfaceSelection,
+  TextSelection,
+} from '@blocksuite/std';
+import {
+  GfxControllerIdentifier,
+  type GfxModel,
+  type GfxPrimitiveElementModel,
+} from '@blocksuite/std/gfx';
+import type { BaseSelection, BlockModel } from '@blocksuite/store';
+import { signal } from '@preact/signals-core';
+
+import { getSelectedBlocksCommand } from '../../commands';
+import { ImageSelection } from '../../selection';
+import { matchModels } from '../../utils';
+import { type CommentId, CommentProviderIdentifier } from './comment-provider';
+import { findCommentedBlocks, findCommentedElements } from './utils';
+
+export class BlockElementCommentManager extends LifeCycleWatcher {
+  static override key = 'block-element-comment-manager';
+
+  private readonly _highlightedCommentId$ = signal<CommentId | null>(null);
+
+  private readonly _disposables = new DisposableGroup();
+
+  private get _provider() {
+    return this.std.getOptional(CommentProviderIdentifier);
+  }
+
+  isBlockCommentHighlighted(
+    block: BlockModel<{ comments?: Record<CommentId, boolean> }>
+  ) {
+    const comments = block.props.comments;
+    if (!comments) return false;
+    return (
+      this._highlightedCommentId$.value !== null &&
+      Object.keys(comments).includes(this._highlightedCommentId$.value)
+    );
+  }
+
+  isElementCommentHighlighted(element: GfxPrimitiveElementModel) {
+    const comments = element.comments;
+    if (!comments) return false;
+    return (
+      this._highlightedCommentId$.value !== null &&
+      Object.keys(comments).includes(this._highlightedCommentId$.value)
+    );
+  }
+
+  override mounted() {
+    const provider = this._provider;
+    if (!provider) return;
+
+    this._disposables.add(provider.onCommentAdded(this._handleAddComment));
+    this._disposables.add(
+      provider.onCommentDeleted(this.handleDeleteAndResolve)
+    );
+    this._disposables.add(
+      provider.onCommentResolved(this.handleDeleteAndResolve)
+    );
+    this._disposables.add(
+      provider.onCommentHighlighted(this._handleHighlightComment)
+    );
+  }
+
+  override unmounted() {
+    this._disposables.dispose();
+  }
+
+  private readonly _handleAddComment = (
+    id: CommentId,
+    selections: BaseSelection[]
+  ) => {
+    // get blocks from text range that some no-text blocks are selected such as image, bookmark, etc.
+    const noTextBlocksFromTextRange = selections
+      .filter((s): s is TextSelection => s.is(TextSelection))
+      .flatMap(s => {
+        const [_, { selectedBlocks }] = this.std.command.exec(
+          getSelectedBlocksCommand,
+          {
+            textSelection: s,
+          }
+        );
+        if (!selectedBlocks) return [];
+        return selectedBlocks.map(b => b.model).filter(m => !m.text);
+      });
+
+    const blocksFromBlockSelection = selections
+      .filter(s => s instanceof BlockSelection || s instanceof ImageSelection)
+      .map(({ blockId }) => this.std.store.getModelById(blockId))
+      .filter(
+        (m): m is BlockModel =>
+          m !== null && !matchModels(m, [DividerBlockModel])
+      );
+
+    const needCommentBlocks = [
+      ...noTextBlocksFromTextRange,
+      ...blocksFromBlockSelection,
+    ];
+
+    if (needCommentBlocks.length !== 0) {
+      this.std.store.withoutTransact(() => {
+        needCommentBlocks.forEach(block => {
+          const comments = (
+            'comments' in block.props &&
+            typeof block.props.comments === 'object' &&
+            block.props.comments !== null
+              ? block.props.comments
+              : {}
+          ) as Record<CommentId, boolean>;
+
+          this.std.store.updateBlock(block, {
+            comments: { [id]: true, ...comments },
+          });
+        });
+      });
+    }
+
+    const gfx = this.std.get(GfxControllerIdentifier);
+    const elementsFromSurfaceSelection = selections
+      .filter(s => s instanceof SurfaceSelection)
+      .flatMap(({ blockId, elements }) => {
+        if (blockId !== gfx.surface?.id) return [];
+        return elements
+          .map(id => gfx.getElementById<GfxModel>(id))
+          .filter(m => m !== null);
+      });
+    if (elementsFromSurfaceSelection.length !== 0) {
+      this.std.store.withoutTransact(() => {
+        elementsFromSurfaceSelection.forEach(element => {
+          const comments =
+            'comments' in element &&
+            typeof element.comments === 'object' &&
+            element.comments !== null
+              ? element.comments
+              : {};
+
+          gfx.updateElement(element, {
+            comments: { [id]: true, ...comments },
+          });
+        });
+      });
+    }
+  };
+
+  readonly handleDeleteAndResolve = (id: CommentId) => {
+    const commentedBlocks = findCommentedBlocks(this.std.store, id);
+    this.std.store.withoutTransact(() => {
+      commentedBlocks.forEach(block => {
+        delete block.props.comments[id];
+      });
+    });
+
+    const commentedElements = findCommentedElements(this.std.store, id);
+    this.std.store.withoutTransact(() => {
+      commentedElements.forEach(element => {
+        delete element.comments[id];
+      });
+    });
+  };
+
+  private readonly _handleHighlightComment = (id: CommentId | null) => {
+    this._highlightedCommentId$.value = id;
+  };
+}

blocksuite/affine/shared/src/services/comment-service/index.ts

@@ -1,3 +1,3 @@
-export * from './block-comment-manager';
+export * from './block-element-comment-manager';
 export * from './comment-provider';
 export * from './utils';

blocksuite/affine/shared/src/services/comment-service/utils.ts

@@ -1,6 +1,10 @@
 import { CommentIcon } from '@blocksuite/icons/lit';
-import { BlockSelection } from '@blocksuite/std';
+import { BlockSelection, SurfaceSelection } from '@blocksuite/std';
-import type { BlockModel, Store } from '@blocksuite/store';
+import type {
+  GfxPrimitiveElementModel,
+  SurfaceBlockModel,
+} from '@blocksuite/std/gfx';
+import { BlockModel, type Store } from '@blocksuite/store';
 
 import type { ToolbarAction } from '../toolbar-service';
 import { type CommentId, CommentProviderIdentifier } from './comment-provider';
@@ -22,6 +26,31 @@ export function findCommentedBlocks(store: Store, commentId: CommentId) {
   });
 }
 
+export function findAllCommentedElements(store: Store) {
+  type CommentedElement = GfxPrimitiveElementModel & {
+    comments: Record<CommentId, boolean>;
+  };
+  const surface = store.getModelsByFlavour('affine:surface')[0] as
+    | SurfaceBlockModel
+    | undefined;
+  if (!surface) return [];
+
+  return surface.elementModels.filter(
+    (element): element is CommentedElement => {
+      return (
+        element.comments !== undefined &&
+        Object.keys(element.comments).length > 0
+      );
+    }
+  );
+}
+
+export function findCommentedElements(store: Store, commentId: CommentId) {
+  return findAllCommentedElements(store).filter(element => {
+    return element.comments[commentId];
+  });
+}
+
 export const blockCommentToolbarButton: Omit<ToolbarAction, 'id'> = {
   tooltip: 'Comment',
   when: ({ std }) => !!std.getOptional(CommentProviderIdentifier),
@@ -29,22 +58,26 @@ export const blockCommentToolbarButton: Omit<ToolbarAction, 'id'> = {
   run: ctx => {
     const commentProvider = ctx.std.getOptional(CommentProviderIdentifier);
     if (!commentProvider) return;
-    const selections = ctx.selection.value;
 
+    const selections = ctx.selection.value;
     const model = ctx.getCurrentModel();
 
-    if (selections.length > 1) {
+    // may be hover on a block or element, in this case
+    // the selection is empty, so we need to get the current model
+    if (model && selections.length === 0) {
+      if (model instanceof BlockModel) {
+        commentProvider.addComment([
+          new BlockSelection({
+            blockId: model.id,
+          }),
+        ]);
+      } else if (ctx.gfx.surface?.id) {
+        commentProvider.addComment([
+          new SurfaceSelection(ctx.gfx.surface.id, [model.id], false),
+        ]);
+      }
+    } else if (selections.length > 0) {
       commentProvider.addComment(selections);
-    } else if (model) {
-      commentProvider.addComment([
-        new BlockSelection({
-          blockId: model.id,
-        }),
-      ]);
-    } else if (selections.length === 1) {
-      commentProvider.addComment(selections);
-    } else {
-      return;
     }
   },
 };

blocksuite/framework/std/src/gfx/model/surface/element-model.ts

@@ -52,6 +52,7 @@ export type BaseElementProps = {
   index: string;
   seed: number;
   lockedBySelf?: boolean;
+  comments?: Record<string, boolean>;
 };
 
 export type SerializedElement = Record<string, unknown> & {
@@ -60,6 +61,7 @@ export type SerializedElement = Record<string, unknown> & {
   id: string;
   index: string;
   lockedBySelf?: boolean;
+  comments?: Record<string, boolean>;
   props: Record<string, unknown>;
 };
 export abstract class GfxPrimitiveElementModel<
@@ -372,6 +374,9 @@ export abstract class GfxPrimitiveElementModel<
 
   @field()
   accessor seed!: number;
+
+  @field()
+  accessor comments: Record<string, boolean> | undefined = undefined;
 }
 
 export abstract class GfxGroupLikeElementModel<

blocksuite/playground/apps/starter/utils/extensions.ts

@@ -33,6 +33,7 @@ export function getTestCommonExtensions(
         di.override(DocModeProvider, mockDocModeService(editor));
       },
     },
+    // CommentProviderExtension(mockCommentProvider()),
   ];
 }
 

packages/frontend/core/src/blocksuite/view-extensions/comment/comment-provider.ts

@@ -5,8 +5,18 @@ import { CommentProviderIdentifier } from '@blocksuite/affine/shared/services';
 import type { BlockStdScope } from '@blocksuite/affine/std';
 import { StdIdentifier } from '@blocksuite/affine/std';
 import type { BaseSelection, ExtensionType } from '@blocksuite/affine/store';
+import { ImageSelection } from '@blocksuite/affine-shared/selection';
 import { type Container } from '@blocksuite/global/di';
-import { BlockSelection, TextSelection } from '@blocksuite/std';
+import {
+  BlockSelection,
+  SurfaceSelection,
+  TextSelection,
+} from '@blocksuite/std';
+import {
+  GfxBlockElementModel,
+  GfxControllerIdentifier,
+  GfxPrimitiveElementModel,
+} from '@blocksuite/std/gfx';
 import type { FrameworkProvider } from '@toeverything/infra';
 
 import { DocCommentManagerService } from '../../../modules/comment/services/doc-comment-manager';
@@ -21,6 +31,8 @@ function getPreviewFromSelections(
 
   const previews: string[] = [];
 
+  const gfx = std.get(GfxControllerIdentifier);
+
   for (const selection of selections) {
     if (selection instanceof TextSelection) {
       // Extract text from TextSelection
@@ -35,9 +47,23 @@ function getPreviewFromSelections(
         const flavour = block.model.flavour.replace('affine:', '');
         previews.push(`<${flavour}>`);
       }
-    } else if (selection.type === 'image') {
+    } else if (selection instanceof ImageSelection) {
       // Return <"Image"> for ImageSelection
       previews.push('<Image>');
+    } else if (
+      selection instanceof SurfaceSelection &&
+      gfx.surface?.id === selection.blockId
+    ) {
+      selection.elements.forEach(elementId => {
+        const model = gfx.getElementById(elementId);
+        if (model instanceof GfxPrimitiveElementModel) {
+          const flavour = model.type.replace('affine:', '');
+          previews.push(`<${flavour}>`);
+        } else if (model instanceof GfxBlockElementModel) {
+          const flavour = model.flavour.replace('affine:', '');
+          previews.push(`<${flavour}>`);
+        }
+      });
     }
     // Skip other types
   }

packages/frontend/core/src/modules/editor/entities/editor.ts

@@ -9,6 +9,7 @@ import { HighlightSelection } from '@blocksuite/affine/shared/selection';
 import {
   DocModeProvider,
   findCommentedBlocks,
+  findCommentedElements,
 } from '@blocksuite/affine/shared/services';
 import { GfxControllerIdentifier } from '@blocksuite/affine/std/gfx';
 import type { InlineEditor } from '@blocksuite/std/inline';
@@ -231,6 +232,12 @@ export class Editor extends Entity {
         if (blockCommentedBlocks.length > 0) {
           finalId = blockCommentedBlocks[0].id;
           finalKey = 'blockIds';
+        } else {
+          const commentedElements = findCommentedElements(std.store, commentId);
+          if (commentedElements.length > 0) {
+            finalId = commentedElements[0].id;
+            finalKey = 'elementIds';
+          }
         }
       }
       // Workaround: clear selection to avoid comment editor flickering