feat(editor): replace spec provider with extension manager (#11861)

Closes: BS-3273

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

@@ -3,10 +3,9 @@ import {
   DocModeProvider,
   FeatureFlagService,
 } from '@blocksuite/affine/shared/services';
-import { type SpecBuilder } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
-import type { BaseSelection } from '@blocksuite/affine/store';
+import type { BaseSelection, ExtensionType } from '@blocksuite/affine/store';
 import { ArrowDownBigIcon as ArrowDownIcon } from '@blocksuite/icons/lit';
 import { css, html, nothing, type PropertyValues } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
@@ -157,7 +156,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
   accessor updateContext!: (context: Partial<ChatContextValue>) => void;
 
   @property({ attribute: false })
-  accessor previewSpecBuilder!: SpecBuilder;
+  accessor extensions!: ExtensionType[];
 
   @query('.chat-panel-messages-container')
   accessor messagesContainer: HTMLDivElement | null = null;
@@ -271,7 +270,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
                     .isLast=${isLast}
                     .status=${isLast ? status : 'idle'}
                     .error=${isLast ? error : null}
-                    .previewSpecBuilder=${this.previewSpecBuilder}
+                    .extensions=${this.extensions}
                     .getSessionId=${this.getSessionId}
                     .retry=${() => this.retry()}
                   ></chat-message-assistant>`;

packages/frontend/core/src/blocksuite/ai/chat-panel/content/rich-text.ts

@@ -1,7 +1,7 @@
 import { WithDisposable } from '@blocksuite/affine/global/lit';
-import type { SpecBuilder } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
+import type { ExtensionType } from '@blocksuite/affine/store';
 import { html } from 'lit';
 import { property } from 'lit/decorators.js';
 
@@ -18,13 +18,13 @@ export class ChatContentRichText extends WithDisposable(ShadowlessElement) {
   accessor state: 'finished' | 'generating' = 'finished';
 
   @property({ attribute: false })
-  accessor previewSpecBuilder!: SpecBuilder;
+  accessor extensions!: ExtensionType[];
 
   protected override render() {
     const { text, host } = this;
     return html`${createTextRenderer(host, {
       customHeading: true,
-      extensions: this.previewSpecBuilder.value,
+      extensions: this.extensions,
     })(text, this.state)}`;
   }
 }

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

@@ -2,10 +2,9 @@ import './chat-panel-messages';
 
 import type { ContextEmbedStatus } from '@affine/graphql';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
-import type { SpecBuilder } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
-import type { Store } from '@blocksuite/affine/store';
+import type { ExtensionType, Store } from '@blocksuite/affine/store';
 import { HelpIcon } from '@blocksuite/icons/lit';
 import { type Signal, signal } from '@preact/signals-core';
 import { css, html, type PropertyValues } from 'lit';
@@ -208,7 +207,7 @@ export class ChatPanel extends SignalWatcher(
   accessor docDisplayConfig!: DocDisplayConfig;
 
   @property({ attribute: false })
-  accessor previewSpecBuilder!: SpecBuilder;
+  accessor extensions!: ExtensionType[];
 
   @state()
   accessor isLoading = false;
@@ -401,7 +400,7 @@ export class ChatPanel extends SignalWatcher(
         .updateContext=${this.updateContext}
         .host=${this.host}
         .isLoading=${this.isLoading}
-        .previewSpecBuilder=${this.previewSpecBuilder}
+        .extensions=${this.extensions}
       ></chat-panel-messages>
       <ai-chat-composer
         .host=${this.host}

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

@@ -5,6 +5,7 @@ import { WithDisposable } from '@blocksuite/affine/global/lit';
 import { isInsidePageEditor } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
+import type { ExtensionType } from '@blocksuite/affine/store';
 import { css, html, nothing } from 'lit';
 import { property } from 'lit/decorators.js';
 
@@ -44,7 +45,7 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
   accessor error: AIError | null = null;
 
   @property({ attribute: false })
-  accessor previewSpecBuilder: any;
+  accessor extensions!: ExtensionType[];
 
   @property({ attribute: false })
   accessor getSessionId!: () => Promise<string | undefined>;
@@ -90,7 +91,7 @@ export class ChatMessageAssistant extends WithDisposable(ShadowlessElement) {
         .host=${host}
         .text=${item.content}
         .state=${state}
-        .previewSpecBuilder=${this.previewSpecBuilder}
+        .extensions=${this.extensions}
       ></chat-content-rich-text>
       ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
       ${this.renderEditorActions()}

packages/frontend/core/src/blocksuite/ai/components/text-renderer.ts

@@ -43,6 +43,18 @@ import type {
   AffineAIPanelWidgetConfig,
 } from '../widgets/ai-panel/type';
 
+export const CustomPageEditorBlockSpecs: ExtensionType[] = [
+  ...PageEditorBlockSpecs,
+  {
+    setup: di => {
+      di.override(
+        BlockViewIdentifier('affine:page'),
+        () => literal`affine-page-root`
+      );
+    },
+  },
+];
+
 const customHeadingStyles = css`
   .custom-heading {
     .h1 {
@@ -91,18 +103,6 @@ export type TextRendererOptions = {
   testId?: string;
 };
 
-export const CustomPageEditorBlockSpecs: ExtensionType[] = [
-  ...PageEditorBlockSpecs,
-  {
-    setup: di => {
-      di.override(
-        BlockViewIdentifier('affine:page'),
-        () => literal`affine-page-root`
-      );
-    },
-  },
-];
-
 // todo: refactor it for more general purpose usage instead of AI only?
 export class TextRenderer extends WithDisposable(ShadowlessElement) {
   static override styles = css`

packages/frontend/core/src/blocksuite/ai/extensions/ai-code.ts

@@ -7,7 +7,7 @@ import type { ExtensionType } from '@blocksuite/affine/store';
 
 import { setupCodeToolbarAIEntry } from '../entries/code-toolbar/setup-code-toolbar';
 
-class AICodeBlockWatcher extends LifeCycleWatcher {
+export class AICodeBlockWatcher extends LifeCycleWatcher {
   static override key = 'ai-code-block-watcher';
 
   override mounted() {

packages/frontend/core/src/blocksuite/ai/extensions/ai-edgeless-root.ts

@@ -47,7 +47,7 @@ export function createAIEdgelessRootBlockSpec(
   ];
 }
 
-function getAIEdgelessRootWatcher(framework: FrameworkProvider) {
+export function getAIEdgelessRootWatcher(framework: FrameworkProvider) {
   class AIEdgelessRootWatcher extends LifeCycleWatcher {
     static override key = 'ai-edgeless-root-watcher';
 

packages/frontend/core/src/blocksuite/ai/extensions/ai-page-root.ts

@@ -16,7 +16,7 @@ import {
 } from '../widgets/ai-panel/ai-panel';
 import { AiSlashMenuConfigExtension } from './ai-slash-menu';
 
-function getAIPageRootWatcher(framework: FrameworkProvider) {
+export function getAIPageRootWatcher(framework: FrameworkProvider) {
   class AIPageRootWatcher extends LifeCycleWatcher {
     static override key = 'ai-page-root-watcher';
 

packages/frontend/core/src/blocksuite/ai/messages/slides-renderer.ts

@@ -1,6 +1,6 @@
 import { WorkspaceImpl } from '@affine/core/modules/workspace/impls/workspace';
+import { ViewExtensionManagerIdentifier } from '@blocksuite/affine/ext-loader';
 import { WithDisposable } from '@blocksuite/affine/global/lit';
-import { SpecProvider } from '@blocksuite/affine/shared/utils';
 import { BlockStdScope, type EditorHost } from '@blocksuite/affine/std';
 import type { Store } from '@blocksuite/affine/store';
 import { css, html, LitElement, nothing } from 'lit';
@@ -91,6 +91,12 @@ export class AISlidesRenderer extends WithDisposable(LitElement) {
     });
   }
 
+  protected _getExtensions() {
+    return this.host.std
+      .get(ViewExtensionManagerIdentifier)
+      .get('preview-edgeless');
+  }
+
   protected override render() {
     return html`<style>
         .slides-container {
@@ -205,7 +211,7 @@ export class AISlidesRenderer extends WithDisposable(LitElement) {
         >
           ${new BlockStdScope({
             store: this._doc,
-            extensions: SpecProvider._.getSpec('preview:edgeless').value,
+            extensions: this._getExtensions(),
           }).render()}
         </div>
         <div class="mask"></div>

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

@@ -4,12 +4,13 @@ import {
   EdgelessCRUDIdentifier,
   getSurfaceBlock,
 } from '@blocksuite/affine/blocks/surface';
+import { ViewExtensionManagerIdentifier } from '@blocksuite/affine/ext-loader';
 import { ConnectorMode } from '@blocksuite/affine/model';
 import {
   DocModeProvider,
   TelemetryProvider,
 } from '@blocksuite/affine/shared/services';
-import type { Signal, SpecBuilder } from '@blocksuite/affine/shared/utils';
+import type { Signal } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { signal } from '@preact/signals-core';
 import { html, LitElement, nothing } from 'lit';
@@ -437,9 +438,11 @@ export class AIChatBlockPeekView extends LitElement {
 
   override connectedCallback() {
     super.connectedCallback();
-    this._textRendererOptions = {
-      extensions: this.previewSpecBuilder.value,
-    };
+    const extensions = this.host.std
+      .get(ViewExtensionManagerIdentifier)
+      .get('preview-page');
+
+    this._textRendererOptions = { extensions };
     this._historyMessages = this._deserializeHistoryChatMessages(
       this.historyMessagesString
     );
@@ -529,9 +532,6 @@ export class AIChatBlockPeekView extends LitElement {
   @property({ attribute: false })
   accessor host!: EditorHost;
 
-  @property({ attribute: false })
-  accessor previewSpecBuilder!: SpecBuilder;
-
   @property({ attribute: false })
   accessor networkSearchConfig!: AINetworkSearchConfig;
 
@@ -566,7 +566,6 @@ declare global {
 export const AIChatBlockPeekViewTemplate = (
   blockModel: AIChatBlockModel,
   host: EditorHost,
-  previewSpecBuilder: SpecBuilder,
   docDisplayConfig: DocDisplayConfig,
   searchMenuConfig: SearchMenuConfig,
   networkSearchConfig: AINetworkSearchConfig
@@ -574,7 +573,6 @@ export const AIChatBlockPeekViewTemplate = (
   return html`<ai-chat-block-peek-view
     .blockModel=${blockModel}
     .host=${host}
-    .previewSpecBuilder=${previewSpecBuilder}
     .networkSearchConfig=${networkSearchConfig}
     .docDisplayConfig=${docDisplayConfig}
     .searchMenuConfig=${searchMenuConfig}

packages/frontend/core/src/blocksuite/block-suite-editor/bi-directional-link-panel.tsx

@@ -41,11 +41,11 @@ import {
   AffineSharedPageReference,
 } from '../../components/affine/reference-link';
 import { LitTextRenderer } from '../ai/components/text-renderer';
-import { enableEditorExtension } from '../extensions/entry/enable-editor';
 import {
   patchReferenceRenderer,
   type ReferenceReactRenderer,
 } from '../extensions/reference-renderer';
+import { getViewManager } from '../manager/migrating-view';
 import * as styles from './bi-directional-link-panel.css';
 
 const PREFIX = 'bi-directional-link-panel-collapse:';
@@ -163,9 +163,9 @@ const usePreviewExtensions = () => {
   const enableAI = useEnableAI();
 
   const extensions = useMemo(() => {
-    const specs = enableEditorExtension(framework, 'page', enableAI);
-    specs.extend([patchReferenceRenderer(reactToLit, referenceRenderer)]);
-    return specs.value;
+    const manager = getViewManager(framework, enableAI);
+    const specs = manager.get('preview-page');
+    return [...specs, patchReferenceRenderer(reactToLit, referenceRenderer)];
   }, [reactToLit, referenceRenderer, framework, enableAI]);
 
   return [extensions, portals] as const;

packages/frontend/core/src/blocksuite/block-suite-editor/index.ts

@@ -1,10 +1,11 @@
 import { registerAIEffects } from '@affine/core/blocksuite/ai/effects';
 import { editorEffects } from '@affine/core/blocksuite/editors';
-import { effects as bsEffects } from '@blocksuite/affine/effects';
+import type * as EffectType from '@blocksuite/affine/effects';
+
+declare type _GLOBAL_ = typeof EffectType;
 
 import { registerTemplates } from './register-templates';
 
-bsEffects();
 editorEffects();
 registerAIEffects();
 registerTemplates();

packages/frontend/core/src/blocksuite/block-suite-editor/lit-adaper.tsx

@@ -62,7 +62,6 @@ import { EdgelessClipboardAIChatConfig } from '../extensions/edgeless-clipboard'
 import { patchForClipboardInElectron } from '../extensions/electron-clipboard';
 import { enableEditorExtension } from '../extensions/entry/enable-editor';
 import { enableMobileExtension } from '../extensions/entry/enable-mobile';
-import { enablePreviewExtension } from '../extensions/entry/enable-preview';
 import { patchForEdgelessNoteConfig } from '../extensions/note-config';
 import { patchNotificationService } from '../extensions/notification-service';
 import { patchOpenDocExtension } from '../extensions/open-doc';
@@ -140,10 +139,6 @@ const usePatchSpecs = (mode: DocMode) => {
     };
   }, [workspaceService]);
 
-  useMemo(() => {
-    enablePreviewExtension(framework);
-  }, [framework]);
-
   const confirmModal = useConfirmModal();
 
   const enableAI = useEnableAI();
@@ -161,9 +156,9 @@ const usePatchSpecs = (mode: DocMode) => {
   );
 
   const patchedSpecs = useMemo(() => {
-    const builder = enableEditorExtension(framework, mode, enableAI);
+    let extensions = enableEditorExtension(framework, mode, enableAI);
 
-    builder.extend(
+    extensions = extensions.concat(
       [
         patchReferenceRenderer(reactToLit, referenceRenderer),
         patchForEdgelessNoteConfig(framework, reactToLit, insidePeekView),
@@ -190,17 +185,20 @@ const usePatchSpecs = (mode: DocMode) => {
     );
 
     if (enablePDFEmbedPreview) {
-      builder.extend([patchForPDFEmbedView(reactToLit)]);
+      extensions = extensions.concat([patchForPDFEmbedView(reactToLit)]);
     }
 
     if (BUILD_CONFIG.isMobileEdition) {
-      enableMobileExtension(builder, framework);
-    }
-    if (BUILD_CONFIG.isElectron) {
-      builder.extend([patchForClipboardInElectron(framework)].flat());
+      extensions = enableMobileExtension(extensions, framework);
     }
 
-    return builder.value;
+    if (BUILD_CONFIG.isElectron) {
+      extensions = extensions.concat(
+        [patchForClipboardInElectron(framework)].flat()
+      );
+    }
+
+    return extensions;
   }, [
     framework,
     mode,

packages/frontend/core/src/blocksuite/extensions/entry/enable-editor.ts

@@ -1,19 +1,13 @@
-import { enableAIExtension } from '@affine/core/blocksuite/ai';
-import { enableAffineExtension } from '@affine/core/blocksuite/extensions';
-import {
-  type SpecBuilder,
-  SpecProvider,
-} from '@blocksuite/affine/shared/utils';
+import type { ExtensionType } from '@blocksuite/affine/store';
 import { type FrameworkProvider } from '@toeverything/infra';
 
+import { getViewManager } from '../../manager/migrating-view';
+
 export function enableEditorExtension(
   framework: FrameworkProvider,
   mode: 'edgeless' | 'page',
   enableAI: boolean
-): SpecBuilder {
-  const spec = SpecProvider._.getSpec(mode);
-  enableAffineExtension(spec, framework);
-  enableAIExtension(spec, framework, enableAI);
-
-  return spec;
+): ExtensionType[] {
+  const manager = getViewManager(framework, enableAI);
+  return manager.get(mode);
 }

packages/frontend/core/src/blocksuite/extensions/entry/enable-mobile.ts

@@ -11,7 +11,6 @@ import {
   VirtualKeyboardProvider as BSVirtualKeyboardProvider,
   type VirtualKeyboardProviderWithAction,
 } from '@blocksuite/affine/shared/services';
-import type { SpecBuilder } from '@blocksuite/affine/shared/utils';
 import { type BlockStdScope, LifeCycleWatcher } from '@blocksuite/affine/std';
 import type { ExtensionType } from '@blocksuite/affine/store';
 import { SlashMenuExtension } from '@blocksuite/affine/widgets/slash-menu';
@@ -112,16 +111,26 @@ function KeyboardToolbarExtension(framework: FrameworkProvider): ExtensionType {
 }
 
 export function enableMobileExtension(
-  specBuilder: SpecBuilder,
+  extensions: ExtensionType[],
   framework: FrameworkProvider
-): void {
-  specBuilder.omit(codeToolbarWidget);
-  specBuilder.omit(toolbarWidget);
-  specBuilder.omit(SlashMenuExtension);
-  specBuilder.extend([
+): ExtensionType[] {
+  const next = extensions.filter(extension => {
+    if (extension === codeToolbarWidget) {
+      return false;
+    }
+    if (extension === toolbarWidget) {
+      return false;
+    }
+    if (extension === SlashMenuExtension) {
+      return false;
+    }
+    return true;
+  });
+  next.push(
     MobileSpecsPatches,
     KeyboardToolbarExtension(framework),
     mobileParagraphConfig,
-    mobileCodeConfig,
-  ]);
+    mobileCodeConfig
+  );
+  return next;
 }

packages/frontend/core/src/blocksuite/extensions/entry/enable-preview.ts

@@ -27,7 +27,7 @@ import { getFontConfigExtension } from '../font-config';
 import { patchPeekViewService } from '../peek-view-service';
 import { getThemeExtension } from '../theme';
 
-function getPagePreviewThemeExtension(framework: FrameworkProvider) {
+export function getPagePreviewThemeExtension(framework: FrameworkProvider) {
   class AffinePagePreviewThemeExtension
     extends LifeCycleWatcher
     implements ThemeExtension

packages/frontend/core/src/blocksuite/manager/migrating-store.ts

@@ -0,0 +1,26 @@
+import {
+  type StoreExtensionContext,
+  StoreExtensionManager,
+  StoreExtensionProvider,
+} from '@blocksuite/affine/ext-loader';
+import { MigratingStoreExtension } from '@blocksuite/affine/extensions/store';
+
+import { AIChatBlockSchemaExtension } from '../ai/blocks/ai-chat-block/model';
+
+class MigratingAffineStoreExtension extends StoreExtensionProvider {
+  override name = 'affine-store-extensions';
+
+  override setup(context: StoreExtensionContext) {
+    super.setup(context);
+    context.register(AIChatBlockSchemaExtension);
+  }
+}
+
+const manager = new StoreExtensionManager([
+  MigratingAffineStoreExtension,
+  MigratingStoreExtension,
+]);
+
+export function getStoreManager() {
+  return manager;
+}

packages/frontend/core/src/blocksuite/manager/migrating-view.ts

@@ -0,0 +1,144 @@
+import { PeekViewService } from '@affine/core/modules/peek-view';
+import { ParagraphBlockConfigExtension } from '@blocksuite/affine/blocks/paragraph';
+import {
+  type ViewExtensionContext,
+  ViewExtensionManager,
+  ViewExtensionProvider,
+} from '@blocksuite/affine/ext-loader';
+import { MigratingViewExtension } from '@blocksuite/affine/extensions/view';
+import { ToolbarModuleExtension } from '@blocksuite/affine/shared/services';
+import { BlockFlavourIdentifier } from '@blocksuite/affine/std';
+import { FrameworkProvider } from '@toeverything/infra';
+import { z } from 'zod';
+
+import { toolbarAIEntryConfig } from '../ai';
+import { AIChatBlockSpec } from '../ai/blocks';
+import { AITranscriptionBlockSpec } from '../ai/blocks/ai-chat-block/ai-transcription-block';
+import { edgelessToolbarAIEntryConfig } from '../ai/entries/edgeless';
+import { imageToolbarAIEntryConfig } from '../ai/entries/image-toolbar/setup-image-toolbar';
+import { AICodeBlockWatcher } from '../ai/extensions/ai-code';
+import { getAIEdgelessRootWatcher } from '../ai/extensions/ai-edgeless-root';
+import { getAIPageRootWatcher } from '../ai/extensions/ai-page-root';
+import { AiSlashMenuConfigExtension } from '../ai/extensions/ai-slash-menu';
+import { CopilotTool } from '../ai/tool/copilot-tool';
+import { aiPanelWidget } from '../ai/widgets/ai-panel/ai-panel';
+import { edgelessCopilotWidget } from '../ai/widgets/edgeless-copilot';
+import { buildDocDisplayMetaExtension } from '../extensions/display-meta';
+import { getEditorConfigExtension } from '../extensions/editor-config';
+import { getPagePreviewThemeExtension } from '../extensions/entry/enable-preview';
+import { getFontConfigExtension } from '../extensions/font-config';
+import { patchPeekViewService } from '../extensions/peek-view-service';
+import { getTelemetryExtension } from '../extensions/telemetry';
+import { getThemeExtension } from '../extensions/theme';
+
+const optionsSchema = z.object({
+  enableAI: z.boolean().optional(),
+  framework: z.instanceof(FrameworkProvider),
+});
+
+class MigratingAffineViewExtension extends ViewExtensionProvider<
+  z.infer<typeof optionsSchema>
+> {
+  override name = 'affine-view-extensions';
+
+  override schema = optionsSchema;
+
+  override setup(
+    context: ViewExtensionContext,
+    options?: z.infer<typeof optionsSchema>
+  ) {
+    super.setup(context);
+    const { framework, enableAI } = options || {};
+    if (framework) {
+      context.register([
+        getThemeExtension(framework),
+        getFontConfigExtension(),
+        buildDocDisplayMetaExtension(framework),
+      ]);
+      if (context.scope === 'page' || context.scope === 'edgeless') {
+        context.register(getTelemetryExtension());
+        context.register(getEditorConfigExtension(framework));
+      }
+      if (
+        context.scope === 'preview-page' ||
+        context.scope === 'preview-edgeless'
+      ) {
+        context.register(getPagePreviewThemeExtension(framework));
+        context.register(patchPeekViewService(framework.get(PeekViewService)));
+      }
+      if (enableAI) {
+        context.register(AIChatBlockSpec);
+        context.register(AITranscriptionBlockSpec);
+        context.register(
+          [
+            AICodeBlockWatcher,
+            ToolbarModuleExtension({
+              id: BlockFlavourIdentifier('custom:affine:image'),
+              config: imageToolbarAIEntryConfig(),
+            }),
+            ParagraphBlockConfigExtension({
+              getPlaceholder: model => {
+                const placeholders = {
+                  text: "Type '/' for commands, 'space' for AI",
+                  h1: 'Heading 1',
+                  h2: 'Heading 2',
+                  h3: 'Heading 3',
+                  h4: 'Heading 4',
+                  h5: 'Heading 5',
+                  h6: 'Heading 6',
+                  quote: '',
+                };
+                return placeholders[model.props.type];
+              },
+            }),
+          ].flat()
+        );
+        if (context.scope === 'edgeless') {
+          context.register([
+            CopilotTool,
+            aiPanelWidget,
+            edgelessCopilotWidget,
+            getAIEdgelessRootWatcher(framework),
+            // In note
+            ToolbarModuleExtension({
+              id: BlockFlavourIdentifier('custom:affine:note'),
+              config: toolbarAIEntryConfig(),
+            }),
+            ToolbarModuleExtension({
+              id: BlockFlavourIdentifier('custom:affine:surface:*'),
+              config: edgelessToolbarAIEntryConfig(),
+            }),
+            AiSlashMenuConfigExtension(),
+          ]);
+        }
+        if (context.scope === 'page') {
+          context.register([
+            aiPanelWidget,
+            getAIPageRootWatcher(framework),
+            ToolbarModuleExtension({
+              id: BlockFlavourIdentifier('custom:affine:note'),
+              config: toolbarAIEntryConfig(),
+            }),
+            AiSlashMenuConfigExtension(),
+          ]);
+        }
+      }
+    }
+  }
+}
+
+const manager = new ViewExtensionManager([
+  MigratingViewExtension,
+  MigratingAffineViewExtension,
+]);
+
+export function getViewManager(
+  framework: FrameworkProvider,
+  enableAI: boolean
+) {
+  manager.configure(MigratingAffineViewExtension, {
+    framework,
+    enableAI,
+  });
+  return manager;
+}

packages/frontend/core/src/desktop/dialogs/setting/general-setting/editor/edgeless/snapshot.tsx

@@ -1,10 +1,10 @@
 import { Skeleton } from '@affine/component';
+import { getViewManager } from '@affine/core/blocksuite/manager/migrating-view';
 import type { EditorSettingSchema } from '@affine/core/modules/editor-setting';
 import { EditorSettingService } from '@affine/core/modules/editor-setting';
 import { EdgelessCRUDIdentifier } from '@blocksuite/affine/blocks/surface';
 import { Bound } from '@blocksuite/affine/global/gfx';
 import { ViewportElementExtension } from '@blocksuite/affine/shared/services';
-import { SpecProvider } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { BlockStdScope } from '@blocksuite/affine/std';
 import {
@@ -14,7 +14,7 @@ import {
 import type { Block, Store } from '@blocksuite/affine/store';
 import { useFramework } from '@toeverything/infra';
 import { isEqual } from 'lodash-es';
-import { useCallback, useEffect, useRef } from 'react';
+import { useCallback, useEffect, useMemo, useRef } from 'react';
 import { map, pairwise } from 'rxjs';
 
 import {
@@ -55,6 +55,13 @@ export const EdgelessSnapshot = (props: Props) => {
   const framework = useFramework();
   const { editorSetting } = framework.get(EditorSettingService);
 
+  const extensions = useMemo(() => {
+    const manager = getViewManager(framework, false);
+    return manager
+      .get('preview-edgeless')
+      .concat([ViewportElementExtension('.ref-viewport')]);
+  }, [framework]);
+
   const updateElements = useCallback(() => {
     const editorHost = editorHostRef.current;
     const doc = docRef.current;
@@ -76,9 +83,7 @@ export const EdgelessSnapshot = (props: Props) => {
 
     const editorHost = new BlockStdScope({
       store: doc,
-      extensions: SpecProvider._.getSpec('preview:edgeless').extend([
-        ViewportElementExtension('.ref-viewport'),
-      ]).value,
+      extensions,
     }).render();
     docRef.current = doc;
     editorHostRef.current?.remove();
@@ -115,7 +120,7 @@ export const EdgelessSnapshot = (props: Props) => {
 
     // append to dom node
     wrapperRef.current.append(editorHost);
-  }, [docName, firstUpdate, updateElements]);
+  }, [docName, extensions, firstUpdate, updateElements]);
 
   useEffect(() => {
     // eslint-disable-next-line @typescript-eslint/no-floating-promises

packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx

@@ -2,12 +2,10 @@ import { ChatPanel } from '@affine/core/blocksuite/ai';
 import type { AffineEditorContainer } from '@affine/core/blocksuite/block-suite-editor';
 import { useAIChatConfig } from '@affine/core/components/hooks/affine/use-ai-chat-config';
 import { WorkbenchService } from '@affine/core/modules/workbench';
+import { ViewExtensionManagerIdentifier } from '@blocksuite/affine/ext-loader';
 import { RefNodeSlotsProvider } from '@blocksuite/affine/inlines/reference';
 import { DocModeProvider } from '@blocksuite/affine/shared/services';
-import {
-  createSignalFromObservable,
-  SpecProvider,
-} from '@blocksuite/affine/shared/utils';
+import { createSignalFromObservable } from '@blocksuite/affine/shared/utils';
 import { useFramework } from '@toeverything/infra';
 import { forwardRef, useEffect, useRef } from 'react';
 
@@ -69,8 +67,9 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
       chatPanelRef.current.docDisplayConfig = docDisplayConfig;
       chatPanelRef.current.searchMenuConfig = searchMenuConfig;
       chatPanelRef.current.networkSearchConfig = networkSearchConfig;
-      chatPanelRef.current.previewSpecBuilder =
-        SpecProvider._.getSpec('preview:page');
+      chatPanelRef.current.extensions = editor.host.std
+        .get(ViewExtensionManagerIdentifier)
+        .get('preview-page');
 
       containerRef.current?.append(chatPanelRef.current);
     } else {

packages/frontend/core/src/modules/peek-view/view/ai-chat-block-peek-view/index.tsx

@@ -2,7 +2,6 @@ import { toReactNode } from '@affine/component';
 import { AIChatBlockPeekViewTemplate } from '@affine/core/blocksuite/ai';
 import type { AIChatBlockModel } from '@affine/core/blocksuite/ai/blocks/ai-chat-block/model/ai-chat-model';
 import { useAIChatConfig } from '@affine/core/components/hooks/affine/use-ai-chat-config';
-import { SpecProvider } from '@blocksuite/affine/shared/utils';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { useMemo } from 'react';
 
@@ -18,11 +17,9 @@ export const AIChatBlockPeekView = ({
   const { docDisplayConfig, searchMenuConfig, networkSearchConfig } =
     useAIChatConfig();
   return useMemo(() => {
-    const previewSpecBuilder = SpecProvider._.getSpec('preview:page');
     const template = AIChatBlockPeekViewTemplate(
       model,
       host,
-      previewSpecBuilder,
       docDisplayConfig,
       searchMenuConfig,
       networkSearchConfig

packages/frontend/core/src/modules/workspace/impls/doc.ts

@@ -1,4 +1,4 @@
-import { SpecProvider } from '@blocksuite/affine/shared/utils';
+import { getStoreManager } from '@affine/core/blocksuite/manager/migrating-store';
 import {
   AwarenessStore,
   type Doc,
@@ -294,11 +294,9 @@ export class DocImpl implements Doc {
       return this._storeMap.get(key) as Store;
     }
 
-    const storeExtensions = SpecProvider._.getSpec('store');
+    const storeExtensions = getStoreManager().get('store');
     const extensionSet = new Set(
-      storeExtensions.value
-        .concat(extensions ?? [])
-        .concat(this.storeExtensions)
+      storeExtensions.concat(extensions ?? []).concat(this.storeExtensions)
     );
 
     const doc = new Store({