diff --git a/blocksuite/affine/shared/src/utils/spec/spec-builder.ts b/blocksuite/affine/shared/src/utils/spec/spec-builder.ts
index 810288647ed3f..922acbacdd0c6 100644
--- a/blocksuite/affine/shared/src/utils/spec/spec-builder.ts
+++ b/blocksuite/affine/shared/src/utils/spec/spec-builder.ts
@@ -19,9 +19,10 @@ export class SpecBuilder {
     this._value = this._value.filter(extension => extension !== target);
   }
 
-  replace(target: ExtensionType, newExtension: ExtensionType) {
-    this._value = this._value.map(extension =>
-      extension === target ? newExtension : extension
-    );
+  replace(target: ExtensionType[], newExtension: ExtensionType[]) {
+    this._value = [
+      ...this._value.filter(extension => !target.includes(extension)),
+      ...newExtension,
+    ];
   }
 }
diff --git a/blocksuite/blocks/src/_specs/common.ts b/blocksuite/blocks/src/_specs/common.ts
index 7e11e58b56135..1aa20bbffa00e 100644
--- a/blocksuite/blocks/src/_specs/common.ts
+++ b/blocksuite/blocks/src/_specs/common.ts
@@ -54,7 +54,7 @@ export const CommonBlockSpecs: ExtensionType[] = [
 ].flat();
 
 export const PageFirstPartyBlockSpecs: ExtensionType[] = [
-  ...CommonBlockSpecs,
+  CommonBlockSpecs,
   NoteBlockSpec,
   PageSurfaceBlockSpec,
   PageSurfaceRefBlockSpec,
@@ -62,7 +62,7 @@ export const PageFirstPartyBlockSpecs: ExtensionType[] = [
 ].flat();
 
 export const EdgelessFirstPartyBlockSpecs: ExtensionType[] = [
-  ...CommonBlockSpecs,
+  CommonBlockSpecs,
 
   EdgelessNoteBlockSpec,
   EdgelessSurfaceBlockSpec,
diff --git a/blocksuite/blocks/src/_specs/preset/edgeless-specs.ts b/blocksuite/blocks/src/_specs/preset/edgeless-specs.ts
index 7c4c90b662710..3225f68ed87e3 100644
--- a/blocksuite/blocks/src/_specs/preset/edgeless-specs.ts
+++ b/blocksuite/blocks/src/_specs/preset/edgeless-specs.ts
@@ -53,7 +53,7 @@ export const EdgelessBuiltInManager: ExtensionType[] = [
 
 export const EdgelessEditorBlockSpecs: ExtensionType[] = [
   EdgelessRootBlockSpec,
-  ...EdgelessFirstPartyBlockSpecs,
+  EdgelessFirstPartyBlockSpecs,
   EdgelessToolExtension,
   EdgelessBuiltInManager,
 ].flat();
diff --git a/blocksuite/blocks/src/root-block/edgeless/edgeless-root-spec.ts b/blocksuite/blocks/src/root-block/edgeless/edgeless-root-spec.ts
index 9dcc44fe2b128..2b0416d85ba36 100644
--- a/blocksuite/blocks/src/root-block/edgeless/edgeless-root-spec.ts
+++ b/blocksuite/blocks/src/root-block/edgeless/edgeless-root-spec.ts
@@ -1,7 +1,6 @@
 import { FileDropExtension } from '@blocksuite/affine-components/drag-indicator';
 import {
   DNDAPIExtension,
-  DocDisplayMetaService,
   DocModeService,
   EmbedOptionService,
   PageViewportServiceExtension,
@@ -99,7 +98,6 @@ const EdgelessCommonExtension: ExtensionType[] = [
   ToolController,
   DNDAPIExtension,
   PageViewportServiceExtension,
-  DocDisplayMetaService,
   RootBlockAdapterExtensions,
   FileDropExtension,
 ].flat();
diff --git a/blocksuite/blocks/src/root-block/page/page-root-spec.ts b/blocksuite/blocks/src/root-block/page/page-root-spec.ts
index a9414ecb579ab..fb8fdd7b91ef0 100644
--- a/blocksuite/blocks/src/root-block/page/page-root-spec.ts
+++ b/blocksuite/blocks/src/root-block/page/page-root-spec.ts
@@ -1,7 +1,6 @@
 import { FileDropExtension } from '@blocksuite/affine-components/drag-indicator';
 import {
   DNDAPIExtension,
-  DocDisplayMetaService,
   DocModeService,
   EmbedOptionService,
   PageViewportServiceExtension,
@@ -70,7 +69,6 @@ const PageCommonExtension: ExtensionType[] = [
   DocModeService,
   ThemeService,
   EmbedOptionService,
-  DocDisplayMetaService,
   PageViewportServiceExtension,
 ];
 
diff --git a/blocksuite/presets/src/editors/editor-container.ts b/blocksuite/presets/src/editors/editor-container.ts
index e3fc77f4fe3a2..e6c7df1dd21bd 100644
--- a/blocksuite/presets/src/editors/editor-container.ts
+++ b/blocksuite/presets/src/editors/editor-container.ts
@@ -5,6 +5,7 @@ import {
 } from '@blocksuite/block-std';
 import {
   type AbstractEditor,
+  DocDisplayMetaService,
   type DocMode,
   EdgelessEditorBlockSpecs,
   PageEditorBlockSpecs,
@@ -91,12 +92,14 @@ export class AffineEditorContainer
   private readonly _doc = signal<Doc>();
 
   private readonly _edgelessSpecs = signal<ExtensionType[]>(
-    EdgelessEditorBlockSpecs
+    EdgelessEditorBlockSpecs.concat(DocDisplayMetaService)
   );
 
   private readonly _mode = signal<DocMode>('page');
 
-  private readonly _pageSpecs = signal<ExtensionType[]>(PageEditorBlockSpecs);
+  private readonly _pageSpecs = signal<ExtensionType[]>(
+    PageEditorBlockSpecs.concat(DocDisplayMetaService)
+  );
 
   private readonly _specs = computed(() =>
     this._mode.value === 'page'
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts
index d9423f4402c2f..fe9721eb8eb1d 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts
@@ -15,6 +15,7 @@ import {
   DividerBlockSpec,
   EditPropsStore,
   EmbedExtensions,
+  FontLoaderService,
   ImageBlockSpec,
   LatexBlockSpec,
   ListBlockSpec,
@@ -36,6 +37,7 @@ const CommonBlockSpecs: ExtensionType[] = [
   BookmarkBlockSpec,
   AttachmentBlockSpec,
   AdapterFactoryExtensions,
+  FontLoaderService,
 ].flat();
 
 export const DefaultBlockSpecs: ExtensionType[] = [
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts
index c2cd8cb89110e..9b6a1c2642694 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/custom/root-block.ts
@@ -1,11 +1,12 @@
 import {
-  AIEdgelessRootBlockSpec,
-  AIPageRootBlockSpec,
+  AICodeBlockSpec,
+  AIImageBlockSpec,
+  AIParagraphBlockSpec,
 } from '@affine/core/blocksuite/presets/ai';
+import { AIChatBlockSpec } from '@affine/core/blocksuite/presets/blocks';
 import { DocService, DocsService } from '@affine/core/modules/doc';
 import { DocDisplayMetaService } from '@affine/core/modules/doc-display-meta';
 import { EditorSettingService } from '@affine/core/modules/editor-setting';
-import { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import { AppThemeService } from '@affine/core/modules/theme';
 import { mixpanel } from '@affine/track';
 import {
@@ -18,18 +19,17 @@ import type {
   DocDisplayMetaExtension,
   DocDisplayMetaParams,
   RootBlockConfig,
+  SpecBuilder,
   TelemetryEventMap,
   ThemeExtension,
 } from '@blocksuite/affine/blocks';
 import {
+  CodeBlockSpec,
   ColorScheme,
   DocDisplayMetaProvider,
-  EdgelessBuiltInManager,
-  EdgelessRootBlockSpec,
-  EdgelessToolExtension,
   EditorSettingExtension,
-  FontLoaderService,
-  PageRootBlockSpec,
+  ImageBlockSpec,
+  ParagraphBlockSpec,
   TelemetryProvider,
   ThemeExtensionIdentifier,
 } from '@blocksuite/affine/blocks';
@@ -258,36 +258,24 @@ export const extendEdgelessPreviewSpec = (function () {
   };
 })();
 
-export function createPageRootBlockSpec(
-  framework: FrameworkProvider
-): ExtensionType[] {
-  const featureFlagService = framework.get(FeatureFlagService);
-  const enableAI = featureFlagService.flags.enable_ai.value;
-  return [
-    enableAI ? AIPageRootBlockSpec : PageRootBlockSpec,
-    FontLoaderService,
-    getThemeExtension(framework),
-    getFontConfigExtension(),
-    getTelemetryExtension(),
-    getEditorConfigExtension(framework),
-    buildDocDisplayMetaExtension(framework),
-  ].flat();
+export function enableAffineExtension(
+  framework: FrameworkProvider,
+  specBuilder: SpecBuilder
+): void {
+  specBuilder.extend(
+    [
+      getThemeExtension(framework),
+      getFontConfigExtension(),
+      getTelemetryExtension(),
+      getEditorConfigExtension(framework),
+      buildDocDisplayMetaExtension(framework),
+    ].flat()
+  );
 }
 
-export function createEdgelessRootBlockSpec(
-  framework: FrameworkProvider
-): ExtensionType[] {
-  const featureFlagService = framework.get(FeatureFlagService);
-  const enableAI = featureFlagService.flags.enable_ai.value;
-  return [
-    enableAI ? AIEdgelessRootBlockSpec : EdgelessRootBlockSpec,
-    FontLoaderService,
-    getThemeExtension(framework),
-    EdgelessToolExtension,
-    EdgelessBuiltInManager,
-    getFontConfigExtension(),
-    getTelemetryExtension(),
-    getEditorConfigExtension(framework),
-    buildDocDisplayMetaExtension(framework),
-  ].flat();
+export function enableAIExtension(specBuilder: SpecBuilder): void {
+  specBuilder.replace(CodeBlockSpec, AICodeBlockSpec);
+  specBuilder.replace(ImageBlockSpec, AIImageBlockSpec);
+  specBuilder.replace(ParagraphBlockSpec, AIParagraphBlockSpec);
+  specBuilder.extend(AIChatBlockSpec);
 }
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts
index 3133d5b1863f2..ab30c6b41bfcc 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts
@@ -1,36 +1,31 @@
+import { AIEdgelessRootBlockSpec } from '@affine/core/blocksuite/presets/ai';
 import { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import { builtInTemplates as builtInEdgelessTemplates } from '@affine/templates/edgeless';
 import { builtInTemplates as builtInStickersTemplates } from '@affine/templates/stickers';
 import type { ExtensionType } from '@blocksuite/affine/block-std';
 import type { TemplateManager } from '@blocksuite/affine/blocks';
 import {
-  EdgelessNoteBlockSpec,
-  EdgelessSurfaceBlockSpec,
-  EdgelessSurfaceRefBlockSpec,
+  EdgelessRootBlockSpec,
   EdgelessTemplatePanel,
-  EdgelessTextBlockSpec,
-  FrameBlockSpec,
+  SpecProvider,
 } from '@blocksuite/affine/blocks';
 import { type FrameworkProvider } from '@toeverything/infra';
 
-import { AIBlockSpecs, DefaultBlockSpecs } from './common';
-import { createEdgelessRootBlockSpec } from './custom/root-block';
+import { enableAffineExtension, enableAIExtension } from './custom/root-block';
 
 export function createEdgelessModeSpecs(
   framework: FrameworkProvider
 ): ExtensionType[] {
   const featureFlagService = framework.get(FeatureFlagService);
   const enableAI = featureFlagService.flags.enable_ai.value;
-  return [
-    ...(enableAI ? AIBlockSpecs : DefaultBlockSpecs),
-    EdgelessSurfaceBlockSpec,
-    EdgelessSurfaceRefBlockSpec,
-    FrameBlockSpec,
-    EdgelessTextBlockSpec,
-    EdgelessNoteBlockSpec,
-    // special
-    createEdgelessRootBlockSpec(framework),
-  ].flat();
+  const edgelessSpec = SpecProvider.getInstance().getSpec('edgeless');
+  enableAffineExtension(framework, edgelessSpec);
+  if (enableAI) {
+    enableAIExtension(edgelessSpec);
+    edgelessSpec.replace(EdgelessRootBlockSpec, AIEdgelessRootBlockSpec);
+  }
+
+  return edgelessSpec.value;
 }
 
 export function effects() {
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts
index aee6a9f6c1ab7..3a50712468125 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts
@@ -1,26 +1,22 @@
+import { AIPageRootBlockSpec } from '@affine/core/blocksuite/presets/ai';
 import { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import type { ExtensionType } from '@blocksuite/affine/block-std';
-import {
-  NoteBlockSpec,
-  PageSurfaceBlockSpec,
-  PageSurfaceRefBlockSpec,
-} from '@blocksuite/affine/blocks';
+import { PageRootBlockSpec, SpecProvider } from '@blocksuite/affine/blocks';
 import { type FrameworkProvider } from '@toeverything/infra';
 
-import { AIBlockSpecs, DefaultBlockSpecs } from './common';
-import { createPageRootBlockSpec } from './custom/root-block';
+import { enableAffineExtension, enableAIExtension } from './custom/root-block';
 
 export function createPageModeSpecs(
   framework: FrameworkProvider
 ): ExtensionType[] {
   const featureFlagService = framework.get(FeatureFlagService);
   const enableAI = featureFlagService.flags.enable_ai.value;
-  return [
-    ...(enableAI ? AIBlockSpecs : DefaultBlockSpecs),
-    PageSurfaceBlockSpec,
-    PageSurfaceRefBlockSpec,
-    NoteBlockSpec,
-    // special
-    createPageRootBlockSpec(framework),
-  ].flat();
+  const provider = SpecProvider.getInstance();
+  const pageSpec = provider.getSpec('page');
+  enableAffineExtension(framework, pageSpec);
+  if (enableAI) {
+    enableAIExtension(pageSpec);
+    pageSpec.replace(PageRootBlockSpec, AIPageRootBlockSpec);
+  }
+  return pageSpec.value;
 }