feat(editor): selection as store extension (#9605)

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

@@ -3,7 +3,6 @@ import {
   CommandExtension,
   FlavourExtension,
 } from '@blocksuite/block-std';
-import { DatabaseSelectionExtension } from '@blocksuite/data-view';
 import type { ExtensionType } from '@blocksuite/store';
 import { literal } from 'lit/static-html.js';
 
@@ -16,6 +15,5 @@ export const DatabaseBlockSpec: ExtensionType[] = [
   DatabaseBlockService,
   CommandExtension(commands),
   BlockViewExtension('affine:database', literal`affine-database`),
-  DatabaseSelectionExtension,
   DatabaseBlockAdapterExtensions,
 ].flat();

blocksuite/affine/block-image/src/components/page-image-block.ts

@@ -1,11 +1,12 @@
 import { ImageSelection } from '@blocksuite/affine-shared/selection';
-import type { BaseSelection, UIEventStateContext } from '@blocksuite/block-std';
+import type { UIEventStateContext } from '@blocksuite/block-std';
 import {
   BlockSelection,
   ShadowlessElement,
   TextSelection,
 } from '@blocksuite/block-std';
 import { WithDisposable } from '@blocksuite/global/utils';
+import type { BaseSelection } from '@blocksuite/store';
 import { css, html, type PropertyValues } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';

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

@@ -1,4 +1,3 @@
-import { ImageSelectionExtension } from '@blocksuite/affine-shared/selection';
 import {
   BlockViewExtension,
   CommandExtension,
@@ -29,6 +28,5 @@ export const ImageBlockSpec: ExtensionType[] = [
     imageToolbar: literal`affine-image-toolbar-widget`,
   }),
   ImageDropOption,
-  ImageSelectionExtension,
   ImageBlockAdapterExtensions,
 ].flat();

blocksuite/affine/block-list/src/list-block.ts

@@ -14,13 +14,14 @@ import {
 } from '@blocksuite/affine-shared/consts';
 import { DocModeProvider } from '@blocksuite/affine-shared/services';
 import { getViewportElement } from '@blocksuite/affine-shared/utils';
-import type { BaseSelection, BlockComponent } from '@blocksuite/block-std';
+import type { BlockComponent } from '@blocksuite/block-std';
 import {
   BlockSelection,
   getInlineRangeProvider,
   TextSelection,
 } from '@blocksuite/block-std';
 import type { InlineRangeProvider } from '@blocksuite/inline';
+import type { BaseSelection } from '@blocksuite/store';
 import { effect } from '@preact/signals-core';
 import { html, nothing, type TemplateResult } from 'lit';
 import { query, state } from 'lit/decorators.js';

blocksuite/affine/block-note/src/note-service.ts

@@ -2,7 +2,6 @@ import { textConversionConfigs } from '@blocksuite/affine-components/rich-text';
 import { NoteBlockSchema } from '@blocksuite/affine-model';
 import { matchFlavours } from '@blocksuite/affine-shared/utils';
 import {
-  type BaseSelection,
   type BlockComponent,
   BlockSelection,
   BlockService,
@@ -11,7 +10,7 @@ import {
   type UIEventHandler,
   type UIEventStateContext,
 } from '@blocksuite/block-std';
-import type { BlockModel } from '@blocksuite/store';
+import type { BaseSelection, BlockModel } from '@blocksuite/store';
 
 import { moveBlockConfigs } from './move-block';
 import { quickActionConfig } from './quick-action';

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

@@ -26,7 +26,6 @@ import {
   SpecProvider,
 } from '@blocksuite/affine-shared/utils';
 import {
-  type BaseSelection,
   BlockComponent,
   BlockSelection,
   BlockServiceWatcher,
@@ -47,7 +46,7 @@ import {
   DisposableGroup,
   type SerializedXYWH,
 } from '@blocksuite/global/utils';
-import { type Store } from '@blocksuite/store';
+import type { BaseSelection, Store } from '@blocksuite/store';
 import { css, html, nothing, type TemplateResult } from 'lit';
 import { query, state } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';

blocksuite/affine/block-surface/src/surface-spec.ts

@@ -1,4 +1,3 @@
-import { HighlightSelectionExtension } from '@blocksuite/affine-shared/selection';
 import {
   BlockViewExtension,
   CommandExtension,
@@ -23,7 +22,6 @@ const CommonSurfaceBlockSpec: ExtensionType[] = [
   FlavourExtension('affine:surface'),
   SurfaceBlockService,
   CommandExtension(commands),
-  HighlightSelectionExtension,
   MindMapView,
   EdgelessCRUDExtension,
   EdgelessLegacySlotExtension,

blocksuite/affine/data-view/src/core/common/selection-schema.ts

@@ -1,4 +1,4 @@
-import { BaseSelection, SelectionExtension } from '@blocksuite/block-std';
+import { BaseSelection, SelectionExtension } from '@blocksuite/store';
 import { z } from 'zod';
 
 import type { DataViewSelection, GetDataViewSelection } from '../types.js';

blocksuite/affine/shared/src/selection/hightlight.ts

@@ -2,7 +2,7 @@ import {
   type ReferenceParams,
   ReferenceParamsSchema,
 } from '@blocksuite/affine-model';
-import { BaseSelection, SelectionExtension } from '@blocksuite/block-std';
+import { BaseSelection, SelectionExtension } from '@blocksuite/store';
 
 export class HighlightSelection extends BaseSelection {
   static override group = 'scene';

blocksuite/affine/shared/src/selection/image.ts

@@ -1,4 +1,4 @@
-import { BaseSelection, SelectionExtension } from '@blocksuite/block-std';
+import { BaseSelection, SelectionExtension } from '@blocksuite/store';
 import z from 'zod';
 
 const ImageSelectionSchema = z.object({

blocksuite/affine/widget-drag-handle/src/utils.ts

@@ -9,13 +9,9 @@ import {
   getClosestBlockComponentByPoint,
   matchFlavours,
 } from '@blocksuite/affine-shared/utils';
-import type {
-  BaseSelection,
-  BlockComponent,
-  EditorHost,
-} from '@blocksuite/block-std';
+import type { BlockComponent, EditorHost } from '@blocksuite/block-std';
 import { Point, Rect } from '@blocksuite/global/utils';
-import type { BlockModel } from '@blocksuite/store';
+import type { BaseSelection, BlockModel } from '@blocksuite/store';
 
 import {
   DRAG_HANDLE_CONTAINER_HEIGHT,

blocksuite/affine/widget-remote-selection/src/doc/doc-remote-selection.ts

@@ -1,12 +1,11 @@
 import { matchFlavours } from '@blocksuite/affine-shared/utils';
 import {
-  type BaseSelection,
   BlockSelection,
   TextSelection,
   WidgetComponent,
 } from '@blocksuite/block-std';
 import { throttle } from '@blocksuite/global/utils';
-import type { UserInfo } from '@blocksuite/store';
+import type { BaseSelection, UserInfo } from '@blocksuite/store';
 import { computed, effect } from '@preact/signals-core';
 import { css, html, nothing, type PropertyValues } from 'lit';
 import { state } from 'lit/decorators.js';

blocksuite/blocks/src/_specs/common.ts

@@ -27,6 +27,10 @@ import {
   RefNodeSlotsExtension,
   RichTextExtensions,
 } from '@blocksuite/affine-components/rich-text';
+import {
+  HighlightSelectionExtension,
+  ImageSelectionExtension,
+} from '@blocksuite/affine-shared/selection';
 import {
   DefaultOpenDocExtension,
   DocDisplayMetaService,
@@ -34,6 +38,13 @@ import {
   FeatureFlagService,
   FontLoaderService,
 } from '@blocksuite/affine-shared/services';
+import {
+  BlockSelectionExtension,
+  CursorSelectionExtension,
+  SurfaceSelectionExtension,
+  TextSelectionExtension,
+} from '@blocksuite/block-std';
+import { DatabaseSelectionExtension } from '@blocksuite/data-view';
 import type { ExtensionType } from '@blocksuite/store';
 
 import { AdapterFactoryExtensions } from '../_common/adapters/extension.js';
@@ -77,4 +88,13 @@ export const EdgelessFirstPartyBlockSpecs: ExtensionType[] = [
   FontLoaderService,
 ].flat();
 
-export const StoreExtensions: ExtensionType[] = [FeatureFlagService];
+export const StoreExtensions: ExtensionType[] = [
+  FeatureFlagService,
+  BlockSelectionExtension,
+  TextSelectionExtension,
+  SurfaceSelectionExtension,
+  CursorSelectionExtension,
+  HighlightSelectionExtension,
+  ImageSelectionExtension,
+  DatabaseSelectionExtension,
+];

blocksuite/blocks/src/root-block/widgets/ai-panel/ai-panel.ts

@@ -7,9 +7,9 @@ import {
   getPageRootByElement,
   stopPropagation,
 } from '@blocksuite/affine-shared/utils';
-import type { BaseSelection } from '@blocksuite/block-std';
 import { WidgetComponent } from '@blocksuite/block-std';
 import { assertExists } from '@blocksuite/global/utils';
+import type { BaseSelection } from '@blocksuite/store';
 import {
   autoPlacement,
   autoUpdate,

blocksuite/blocks/src/root-block/widgets/format-bar/format-bar.ts

@@ -9,7 +9,6 @@ import {
 import type { AffineTextAttributes } from '@blocksuite/affine-shared/types';
 import { matchFlavours } from '@blocksuite/affine-shared/utils';
 import {
-  type BaseSelection,
   type BlockComponent,
   BlockSelection,
   CursorSelection,
@@ -22,6 +21,7 @@ import {
   DisposableGroup,
   nextTick,
 } from '@blocksuite/global/utils';
+import type { BaseSelection } from '@blocksuite/store';
 import {
   autoUpdate,
   computePosition,

blocksuite/framework/block-std/src/extension/index.ts

@@ -4,7 +4,6 @@ export * from './config.js';
 export * from './flavour.js';
 export * from './keymap.js';
 export * from './lifecycle-watcher.js';
-export * from './selection.js';
 export * from './service.js';
 export * from './service-watcher.js';
 export * from './widget-view-map.js';

blocksuite/framework/block-std/src/extension/selection.ts

@@ -1,14 +0,0 @@
-import type { ExtensionType } from '@blocksuite/store';
-
-import { SelectionIdentifier } from '../identifier.js';
-import type { SelectionConstructor } from '../selection/index.js';
-
-export function SelectionExtension(
-  selectionCtor: SelectionConstructor
-): ExtensionType {
-  return {
-    setup: di => {
-      di.addImpl(SelectionIdentifier(selectionCtor.type), () => selectionCtor);
-    },
-  };
-}

blocksuite/framework/block-std/src/identifier.ts

@@ -4,7 +4,6 @@ import type { Command } from './command/index.js';
 import type { EventOptions, UIEventHandler } from './event/index.js';
 import type { BlockService, LifeCycleWatcher } from './extension/index.js';
 import type { BlockStdScope } from './scope/index.js';
-import type { SelectionConstructor } from './selection/index.js';
 import type { BlockViewType, WidgetViewMapType } from './spec/type.js';
 
 export const BlockServiceIdentifier =
@@ -33,6 +32,3 @@ export const KeymapIdentifier = createIdentifier<{
   getter: (std: BlockStdScope) => Record<string, UIEventHandler>;
   options?: EventOptions;
 }>('Keymap');
-
-export const SelectionIdentifier =
-  createIdentifier<SelectionConstructor>('Selection');

blocksuite/framework/block-std/src/range/inline-range-provider.ts

@@ -84,18 +84,21 @@ export const getInlineRangeProvider: (
     }
   };
   const inlineRange$: InlineRangeProvider['inlineRange$'] = signal(null);
-  selectionManager.slots.changed.on(selections => {
-    const textSelection = selections.find(s => s.type === 'text') as
-      | TextSelection
-      | undefined;
-    const range = rangeManager.value;
-    if (!range || !textSelection) {
-      inlineRange$.value = null;
-      return;
-    }
-    const inlineRange = calculateInlineRange(range, textSelection);
-    inlineRange$.value = inlineRange;
-  });
+
+  editorHost.disposables.add(
+    selectionManager.slots.changed.on(selections => {
+      const textSelection = selections.find(s => s.type === 'text') as
+        | TextSelection
+        | undefined;
+      const range = rangeManager.value;
+      if (!range || !textSelection) {
+        inlineRange$.value = null;
+        return;
+      }
+      const inlineRange = calculateInlineRange(range, textSelection);
+      inlineRange$.value = inlineRange;
+    })
+  );
 
   return {
     setInlineRange,

blocksuite/framework/block-std/src/range/range-binding.ts

@@ -1,7 +1,7 @@
 import { throttle } from '@blocksuite/global/utils';
-import type { BlockModel } from '@blocksuite/store';
+import type { BaseSelection, BlockModel } from '@blocksuite/store';
 
-import { type BaseSelection, TextSelection } from '../selection/index.js';
+import { TextSelection } from '../selection/index.js';
 import type { BlockComponent } from '../view/element/block-component.js';
 import { BLOCK_ID_ATTR } from '../view/index.js';
 import { RANGE_SYNC_EXCLUDE_ATTR } from './consts.js';
@@ -247,6 +247,9 @@ export class RangeBinding {
   };
 
   private readonly _onStdSelectionChanged = (selections: BaseSelection[]) => {
+    const closestHost = document.activeElement?.closest('editor-host');
+    if (closestHost && closestHost !== this.host) return;
+
     const text =
       selections.find((selection): selection is TextSelection =>
         selection.is(TextSelection)

blocksuite/framework/block-std/src/scope/block-std-scope.ts

@@ -5,6 +5,7 @@ import {
   Job,
   type JobMiddleware,
   type Store,
+  StoreSelectionExtension,
 } from '@blocksuite/store';
 
 import { Clipboard } from '../clipboard/index.js';
@@ -23,13 +24,6 @@ import {
   StdIdentifier,
 } from '../identifier.js';
 import { RangeManager } from '../range/index.js';
-import {
-  BlockSelectionExtension,
-  CursorSelectionExtension,
-  SelectionManager,
-  SurfaceSelectionExtension,
-  TextSelectionExtension,
-} from '../selection/index.js';
 import { ServiceManager } from '../service/index.js';
 import { EditorHost } from '../view/element/index.js';
 import { ViewStore } from '../view/view-store.js';
@@ -43,15 +37,10 @@ const internalExtensions = [
   ServiceManager,
   CommandManager,
   UIEventDispatcher,
-  SelectionManager,
   RangeManager,
   ViewStore,
   Clipboard,
   GfxController,
-  BlockSelectionExtension,
-  TextSelectionExtension,
-  SurfaceSelectionExtension,
-  CursorSelectionExtension,
   GfxSelectionManager,
   SurfaceMiddlewareExtension,
   ViewManager,
@@ -107,7 +96,7 @@ export class BlockStdScope {
   }
 
   get selection() {
-    return this.get(SelectionManager);
+    return this.get(StoreSelectionExtension);
   }
 
   get view() {

blocksuite/framework/block-std/src/selection/index.ts

@@ -1,3 +1 @@
-export * from './base.js';
-export * from './manager.js';
 export * from './variants/index.js';

blocksuite/framework/block-std/src/selection/variants/block.ts

@@ -1,8 +1,6 @@
+import { BaseSelection, SelectionExtension } from '@blocksuite/store';
 import z from 'zod';
 
-import { SelectionExtension } from '../../extension/selection.js';
-import { BaseSelection } from '../base.js';
-
 const BlockSelectionSchema = z.object({
   blockId: z.string(),
 });

blocksuite/framework/block-std/src/selection/variants/cursor.ts

@@ -1,8 +1,6 @@
+import { BaseSelection, SelectionExtension } from '@blocksuite/store';
 import z from 'zod';
 
-import { SelectionExtension } from '../../extension/selection.js';
-import { BaseSelection } from '../base.js';
-
 const CursorSelectionSchema = z.object({
   x: z.number(),
   y: z.number(),

blocksuite/framework/block-std/src/selection/variants/surface.ts

@@ -1,8 +1,6 @@
+import { BaseSelection, SelectionExtension } from '@blocksuite/store';
 import z from 'zod';
 
-import { SelectionExtension } from '../../extension/selection.js';
-import { BaseSelection } from '../base.js';
-
 const SurfaceSelectionSchema = z.object({
   blockId: z.string(),
   elements: z.array(z.string()),

blocksuite/framework/block-std/src/selection/variants/text.ts

@@ -1,8 +1,6 @@
+import { BaseSelection, SelectionExtension } from '@blocksuite/store';
 import z from 'zod';
 
-import { SelectionExtension } from '../../extension/selection.js';
-import { BaseSelection } from '../base.js';
-
 export type TextRangePoint = {
   blockId: string;
   index: number;

blocksuite/framework/block-std/src/view/element/lit-host.ts

@@ -4,7 +4,11 @@ import {
   handleError,
 } from '@blocksuite/global/exceptions';
 import { SignalWatcher, Slot, WithDisposable } from '@blocksuite/global/utils';
-import { type BlockModel, Store } from '@blocksuite/store';
+import {
+  type BlockModel,
+  Store,
+  type StoreSelectionExtension,
+} from '@blocksuite/store';
 import { createContext, provide } from '@lit/context';
 import { css, LitElement, nothing, type TemplateResult } from 'lit';
 import { property } from 'lit/decorators.js';
@@ -16,7 +20,6 @@ import type { UIEventDispatcher } from '../../event/index.js';
 import { WidgetViewMapIdentifier } from '../../identifier.js';
 import type { RangeManager } from '../../range/index.js';
 import type { BlockStdScope } from '../../scope/block-std-scope.js';
-import type { SelectionManager } from '../../selection/index.js';
 import { PropTypes, requiredProperties } from '../decorators/index.js';
 import type { ViewStore } from '../view-store.js';
 import { BLOCK_ID_ATTR, WIDGET_ID_ATTR } from './consts.js';
@@ -114,7 +117,7 @@ export class EditorHost extends SignalWatcher(
     return this.std.range;
   }
 
-  get selection(): SelectionManager {
+  get selection(): StoreSelectionExtension {
     return this.std.selection;
   }
 

blocksuite/framework/store/src/extension/index.ts

@@ -1,2 +1,3 @@
 export * from './extension';
+export * from './selection';
 export * from './store-extension';

blocksuite/framework/store/src/extension/selection/base.ts

@@ -1,6 +1,6 @@
 import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
 
-import type { SelectionConstructor } from './manager';
+import type { SelectionConstructor } from './types';
 
 export type BaseSelectionOptions = {
   blockId: string;

blocksuite/framework/store/src/extension/selection/identifier.ts

@@ -0,0 +1,17 @@
+import { createIdentifier } from '@blocksuite/global/di';
+
+import type { ExtensionType } from '../extension';
+import type { SelectionConstructor } from './types';
+
+export const SelectionIdentifier =
+  createIdentifier<SelectionConstructor>('Selection');
+
+export function SelectionExtension(
+  selectionCtor: SelectionConstructor
+): ExtensionType {
+  return {
+    setup: di => {
+      di.addImpl(SelectionIdentifier(selectionCtor.type), () => selectionCtor);
+    },
+  };
+}

blocksuite/framework/store/src/extension/selection/index.ts

@@ -0,0 +1,4 @@
+export * from './base';
+export * from './identifier';
+export * from './selection-extension';
+export * from './types';

blocksuite/framework/store/src/extension/selection/selection-extension.ts

@@ -1,28 +1,27 @@
 import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
-import { DisposableGroup, Slot } from '@blocksuite/global/utils';
-import { nanoid, type StackItem } from '@blocksuite/store';
+import { Slot } from '@blocksuite/global/utils';
 import { computed, signal } from '@preact/signals-core';
 
-import { LifeCycleWatcher } from '../extension/index.js';
-import { SelectionIdentifier } from '../identifier.js';
-import type { BlockStdScope } from '../scope/index.js';
-import type { BaseSelection } from './base.js';
+import type { Store } from '../../model';
+import { nanoid } from '../../utils/id-generator';
+import type { StackItem } from '../../yjs';
+import { StoreExtension } from '../store-extension';
+import type { BaseSelection } from './base';
+import { SelectionIdentifier } from './identifier';
+import type { SelectionConstructor } from './types';
 
-export interface SelectionConstructor<T extends BaseSelection = BaseSelection> {
-  type: string;
-  group: string;
+export class StoreSelectionExtension extends StoreExtension {
+  static override readonly key = 'selection';
 
-  new (...args: any[]): T;
-  fromJSON(json: Record<string, unknown>): T;
-}
-
-export class SelectionManager extends LifeCycleWatcher {
-  static override readonly key = 'selectionManager';
-
-  private readonly _id: string;
+  private readonly _id = `${this.store.id}:${nanoid()}`;
+  private _selectionConstructors: Record<string, SelectionConstructor> = {};
+  private readonly _selections = signal<BaseSelection[]>([]);
+  private readonly _remoteSelections = signal<Map<number, BaseSelection[]>>(
+    new Map()
+  );
 
   private readonly _itemAdded = (event: { stackItem: StackItem }) => {
-    event.stackItem.meta.set('selection-state', this.value);
+    event.stackItem.meta.set('selection-state', this._selections.value);
   };
 
   private readonly _itemPopped = (event: { stackItem: StackItem }) => {
@@ -43,50 +42,30 @@ export class SelectionManager extends LifeCycleWatcher {
     return ctor.fromJSON(json);
   };
 
-  private readonly _remoteSelections = signal<Map<number, BaseSelection[]>>(
-    new Map()
-  );
-
-  private _selectionConstructors: Record<string, SelectionConstructor> = {};
-
-  private readonly _selections = signal<BaseSelection[]>([]);
-
-  disposables = new DisposableGroup();
-
   slots = {
     changed: new Slot<BaseSelection[]>(),
     remoteChanged: new Slot<Map<number, BaseSelection[]>>(),
   };
 
-  private get _store() {
-    return this.std.workspace.awarenessStore;
-  }
+  constructor(store: Store) {
+    super(store);
 
-  get id() {
-    return this._id;
-  }
+    this.store.provider.getAll(SelectionIdentifier).forEach(ctor => {
+      [ctor].flat().forEach(ctor => {
+        this._selectionConstructors[ctor.type] = ctor;
+      });
+    });
 
-  get remoteSelections() {
-    return this._remoteSelections.value;
-  }
-
-  get value() {
-    return this._selections.value;
-  }
-
-  constructor(std: BlockStdScope) {
-    super(std);
-    this._id = `${this.std.store.id}:${nanoid()}`;
-    this._setupDefaultSelections();
-    this._store.awareness.on(
+    this.store.awarenessStore.awareness.on(
       'change',
       (change: { updated: number[]; added: number[]; removed: number[] }) => {
         const all = change.updated.concat(change.added).concat(change.removed);
-        const localClientID = this._store.awareness.clientID;
+        const localClientID = this.store.awarenessStore.awareness.clientID;
         const exceptLocal = all.filter(id => id !== localClientID);
         const hasLocal = all.includes(localClientID);
         if (hasLocal) {
-          const localSelectionJson = this._store.getLocalSelection(this.id);
+          const localSelectionJson =
+            this.store.awarenessStore.getLocalSelection(this._id);
           const localSelection = localSelectionJson.map(json => {
             return this._jsonToSelection(json);
           });
@@ -96,11 +75,11 @@ export class SelectionManager extends LifeCycleWatcher {
         // Only consider remote selections from other clients
         if (exceptLocal.length > 0) {
           const map = new Map<number, BaseSelection[]>();
-          this._store.getStates().forEach((state, id) => {
-            if (id === this._store.awareness.clientID) return;
+          this.store.awarenessStore.getStates().forEach((state, id) => {
+            if (id === this.store.awarenessStore.awareness.clientID) return;
             // selection id starts with the same block collection id from others clients would be considered as remote selections
             const selection = Object.entries(state.selectionV2)
-              .filter(([key]) => key.startsWith(this.std.store.id))
+              .filter(([key]) => key.startsWith(this.store.id))
               .flatMap(([_, selection]) => selection);
 
             const selections = selection
@@ -122,15 +101,21 @@ export class SelectionManager extends LifeCycleWatcher {
             map.set(id, selections);
           });
           this._remoteSelections.value = map;
+          this.slots.remoteChanged.emit(map);
         }
       }
     );
+
+    this.store.history.on('stack-item-added', this._itemAdded);
+    this.store.history.on('stack-item-popped', this._itemPopped);
   }
 
-  private _setupDefaultSelections() {
-    this.std.provider.getAll(SelectionIdentifier).forEach(ctor => {
-      this.register(ctor);
-    });
+  get value() {
+    return this._selections.value;
+  }
+
+  get remoteSelections() {
+    return this._remoteSelections.value;
   }
 
   clear(types?: string[]) {
@@ -151,9 +136,8 @@ export class SelectionManager extends LifeCycleWatcher {
     return new Type(...args) as InstanceType<T>;
   }
 
-  dispose() {
-    Object.values(this.slots).forEach(slot => slot.dispose());
-    this.disposables.dispose();
+  getGroup(group: string) {
+    return this.value.filter(s => s.group === group);
   }
 
   filter<T extends SelectionConstructor>(type: T) {
@@ -176,43 +160,9 @@ export class SelectionManager extends LifeCycleWatcher {
     );
   }
 
-  fromJSON(json: Record<string, unknown>[]) {
-    const selections = json.map(json => {
-      return this._jsonToSelection(json);
-    });
-    return this.set(selections);
-  }
-
-  getGroup(group: string) {
-    return this.value.filter(s => s.group === group);
-  }
-
-  override mounted() {
-    if (this.disposables.disposed) {
-      this.disposables = new DisposableGroup();
-    }
-    this.std.store.history.on('stack-item-added', this._itemAdded);
-    this.std.store.history.on('stack-item-popped', this._itemPopped);
-    this.disposables.add(
-      this._store.slots.update.on(({ id }) => {
-        if (id === this._store.awareness.clientID) {
-          return;
-        }
-        this.slots.remoteChanged.emit(this.remoteSelections);
-      })
-    );
-  }
-
-  register(ctor: SelectionConstructor | SelectionConstructor[]) {
-    [ctor].flat().forEach(ctor => {
-      this._selectionConstructors[ctor.type] = ctor;
-    });
-    return this;
-  }
-
   set(selections: BaseSelection[]) {
-    this._store.setLocalSelection(
-      this.id,
+    this.store.awarenessStore.setLocalSelection(
+      this._id,
       selections.map(s => s.toJSON())
     );
     this.slots.changed.emit(selections);
@@ -223,16 +173,15 @@ export class SelectionManager extends LifeCycleWatcher {
     this.set([...current, ...selections]);
   }
 
-  override unmounted() {
-    this.std.store.history.off('stack-item-added', this._itemAdded);
-    this.std.store.history.off('stack-item-popped', this._itemPopped);
-    this.slots.changed.dispose();
-    this.disposables.dispose();
-    this.clear();
-  }
-
   update(fn: (currentSelections: BaseSelection[]) => BaseSelection[]) {
     const selections = fn(this.value);
     this.set(selections);
   }
+
+  fromJSON(json: Record<string, unknown>[]) {
+    const selections = json.map(json => {
+      return this._jsonToSelection(json);
+    });
+    return this.set(selections);
+  }
 }

blocksuite/framework/store/src/extension/selection/types.ts

@@ -0,0 +1,9 @@
+import type { BaseSelection } from './base';
+
+export interface SelectionConstructor<T extends BaseSelection = BaseSelection> {
+  type: string;
+  group: string;
+
+  new (...args: any[]): T;
+  fromJSON(json: Record<string, unknown>): T;
+}

blocksuite/framework/store/src/model/store/store.ts

@@ -4,6 +4,7 @@ import { type Disposable, Slot } from '@blocksuite/global/utils';
 import { signal } from '@preact/signals-core';
 
 import type { ExtensionType } from '../../extension/extension.js';
+import { StoreSelectionExtension } from '../../extension/index.js';
 import type { Schema } from '../../schema/index.js';
 import {
   Block,
@@ -27,29 +28,33 @@ export type StoreOptions = {
   extensions?: ExtensionType[];
 };
 
+const internalExtensions = [StoreSelectionExtension];
+
 export class Store {
+  readonly userExtensions: ExtensionType[];
+
   private readonly _provider: ServiceProvider;
 
   private readonly _runQuery = (block: Block) => {
     runQuery(this._query, block);
   };
 
-  protected readonly _doc: Doc;
+  private readonly _doc: Doc;
 
-  protected readonly _blocks = signal<Record<string, Block>>({});
+  private readonly _blocks = signal<Record<string, Block>>({});
 
-  protected readonly _crud: DocCRUD;
+  private readonly _crud: DocCRUD;
 
-  protected readonly _disposeBlockUpdated: Disposable;
+  private readonly _disposeBlockUpdated: Disposable;
 
-  protected readonly _query: Query = {
+  private readonly _query: Query = {
     match: [],
     mode: 'loose',
   };
 
-  protected _readonly = signal(false);
+  private readonly _readonly = signal(false);
 
-  protected readonly _schema: Schema;
+  private readonly _schema: Schema;
 
   readonly slots: Doc['slots'] & {
     /** This is always triggered after `doc.load` is called. */
@@ -295,7 +300,12 @@ export class Store {
     const container = new Container();
     container.addImpl(StoreIdentifier, () => this);
 
+    internalExtensions.forEach(ext => {
+      ext.setup(container);
+    });
+
     const userExtensions = extensions ?? [];
+    this.userExtensions = userExtensions;
     userExtensions.forEach(extension => {
       extension.setup(container);
     });

blocksuite/presets/src/__tests__/utils/setup.ts

@@ -9,8 +9,8 @@ effects();
 import {
   CommunityCanvasTextFonts,
   type DocMode,
-  FeatureFlagService,
   FontConfigExtension,
+  StoreExtensions,
 } from '@blocksuite/blocks';
 import { AffineSchemas } from '@blocksuite/blocks/schemas';
 import { assertExists } from '@blocksuite/global/utils';
@@ -85,7 +85,7 @@ async function createEditor(collection: TestWorkspace, mode: DocMode = 'page') {
 
 export async function setupEditor(mode: DocMode = 'page') {
   const collection = new TestWorkspace(createCollectionOptions());
-  collection.storeExtensions = [FeatureFlagService];
+  collection.storeExtensions = StoreExtensions;
   collection.meta.initialize();
 
   window.collection = collection;

blocksuite/tests-legacy/edgeless/note/note.spec.ts

@@ -329,7 +329,6 @@ test('cursor for active and inactive state', async ({ page }) => {
 
   await switchEditorMode(page);
 
-  await assertTextSelection(page);
   await page.mouse.click(CENTER_X, CENTER_Y);
   await waitNextFrame(page);
   await assertTextSelection(page);

blocksuite/tests-legacy/edgeless/paste-block.spec.ts

@@ -43,6 +43,7 @@ test.describe('pasting blocks', () => {
     await focusRichText(page);
     await initContent(page);
     await switchEditorMode(page);
+    await click(page, { x: 0, y: 0 });
     const box = await getNoteBoundBoxInEdgeless(page, noteId);
     await click(page, {
       x: box.x + 10,

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

@@ -1,4 +1,4 @@
-import type { BaseSelection, EditorHost } from '@blocksuite/affine/block-std';
+import type { EditorHost } from '@blocksuite/affine/block-std';
 import { ShadowlessElement } from '@blocksuite/affine/block-std';
 import {
   type AIError,
@@ -9,6 +9,7 @@ import {
   UnauthorizedError,
 } from '@blocksuite/affine/blocks';
 import { WithDisposable } from '@blocksuite/affine/global/utils';
+import type { BaseSelection } from '@blocksuite/affine/store';
 import { css, html, nothing } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';

tests/affine-local/e2e/links.spec.ts

@@ -62,6 +62,7 @@ test('not allowed to switch to embed view when linking to the same document', as
   await expect(peekViewModel.locator('page-editor')).toBeVisible();
   await page.keyboard.press('Escape');
   await expect(peekViewModel).not.toBeVisible();
+  await page.click('body');
 
   await cardLink.click();
   await cardToolbar.getByLabel('Switch view').click();
@@ -103,6 +104,7 @@ test('not allowed to switch to embed view when linking to block', async ({
   await page.keyboard.press('Escape');
   await expect(peekViewModel).not.toBeVisible();
 
+  await page.click('body');
   await cardLink.click();
 
   await cardToolbar.getByLabel('More').click();
@@ -131,6 +133,7 @@ test('not allowed to switch to embed view when linking to block', async ({
   await page.keyboard.press('Escape');
   await expect(peekViewModel).not.toBeVisible();
 
+  await page.click('body');
   await otherCardLink.click();
   await cardToolbar.getByLabel('Switch view').click();