refactor(editor): optimize ai code structure (#10381)

Let me analyze this diff and provide a clear description of the changes.

This PR introduces several significant changes focused on AI integration and code organization in the AFFiNE codebase:

1. **Enhanced SpecBuilder Functionality** (`blocksuite/affine/shared/src/utils/spec/spec-builder.ts`):
   - Added method chaining by returning `this` from `extend`, `omit`, and `replace` methods
   - Added new utility methods:
     - `hasAll(target: ExtensionType[])`: Checks if all specified extensions exist
     - `hasOneOf(target: ExtensionType[])`: Checks if at least one specified extension exists

2. **AI Extensions Modularization**:
   - Split the large AI-related code into separate modular files under `packages/frontend/core/src/blocksuite/ai/extensions/`:
     - `ai-code.ts`: Code block AI integration
     - `ai-edgeless-root.ts`: Edgeless mode AI features
     - `ai-image.ts`: Image block AI capabilities
     - `ai-page-root.ts`: Page root AI integration
     - `ai-paragraph.ts`: Paragraph block AI features
     - `enable-ai.ts`: Central AI extension enablement logic

3. **Widget Improvements**:
   - Enhanced `AffineAIPanelWidget` and `EdgelessCopilotWidget` with proper widget extensions
   - Moved widget-specific extensions into their respective files
   - Added proper type definitions and component registrations

4. **Code Organization**:
   - Simplified exports in `index.ts`
   - Better separation of concerns between different AI-related components
   - More modular approach to AI feature integration

5. **AI Integration Architecture**:
   - Introduced a new `enableAIExtension` function that handles:
     - Replacing standard blocks with AI-enhanced versions
     - Conditional enabling of AI features based on the current spec configuration
     - Extension of AI chat capabilities

The changes primarily focus on improving code organization, maintainability, and the architecture of AI feature integration in the AFFiNE editor. The modularization will make it easier to maintain and extend AI capabilities across different block types and editor modes.

packages/frontend/core/src/blocksuite/ai/ai-spec.ts

@@ -1,199 +0,0 @@
-import {
-  BlockServiceWatcher,
-  WidgetViewExtension,
-} from '@blocksuite/affine/block-std';
-import {
-  AffineCodeToolbarWidget,
-  AffineFormatBarWidget,
-  AffineImageToolbarWidget,
-  AffineSlashMenuWidget,
-  CodeBlockSpec,
-  EdgelessElementToolbarWidget,
-  EdgelessRootBlockSpec,
-  ImageBlockSpec,
-  PageRootBlockSpec,
-  ParagraphBlockService,
-  ParagraphBlockSpec,
-} from '@blocksuite/affine/blocks';
-import { assertInstanceOf } from '@blocksuite/affine/global/utils';
-import type { ExtensionType } from '@blocksuite/affine/store';
-import type { FrameworkProvider } from '@toeverything/infra';
-import { literal, unsafeStatic } from 'lit/static-html.js';
-
-import { buildAIPanelConfig } from './ai-panel';
-import { setupCodeToolbarAIEntry } from './entries/code-toolbar/setup-code-toolbar';
-import {
-  setupEdgelessCopilot,
-  setupEdgelessElementToolbarAIEntry,
-} from './entries/edgeless/index';
-import { setupFormatBarAIEntry } from './entries/format-bar/setup-format-bar';
-import { setupImageToolbarAIEntry } from './entries/image-toolbar/setup-image-toolbar';
-import { setupSlashMenuAIEntry } from './entries/slash-menu/setup-slash-menu';
-import { setupSpaceAIEntry } from './entries/space/setup-space';
-import { CopilotTool } from './tool/copilot-tool';
-import {
-  AFFINE_AI_PANEL_WIDGET,
-  AffineAIPanelWidget,
-} from './widgets/ai-panel/ai-panel';
-import {
-  AFFINE_EDGELESS_COPILOT_WIDGET,
-  EdgelessCopilotWidget,
-} from './widgets/edgeless-copilot';
-
-function getAIPageRootWatcher(framework: FrameworkProvider) {
-  class AIPageRootWatcher extends BlockServiceWatcher {
-    static override readonly flavour = 'affine:page';
-
-    override mounted() {
-      super.mounted();
-      this.blockService.specSlots.widgetConnected.on(view => {
-        if (view.component instanceof AffineAIPanelWidget) {
-          view.component.style.width = '630px';
-          view.component.config = buildAIPanelConfig(view.component, framework);
-          setupSpaceAIEntry(view.component);
-        }
-
-        if (view.component instanceof AffineFormatBarWidget) {
-          setupFormatBarAIEntry(view.component);
-        }
-
-        if (view.component instanceof AffineSlashMenuWidget) {
-          setupSlashMenuAIEntry(view.component);
-        }
-      });
-    }
-  }
-  return AIPageRootWatcher;
-}
-
-const aiPanelWidget = WidgetViewExtension(
-  'affine:page',
-  AFFINE_AI_PANEL_WIDGET,
-  literal`${unsafeStatic(AFFINE_AI_PANEL_WIDGET)}`
-);
-
-const edgelessCopilotWidget = WidgetViewExtension(
-  'affine:page',
-  AFFINE_EDGELESS_COPILOT_WIDGET,
-  literal`${unsafeStatic(AFFINE_EDGELESS_COPILOT_WIDGET)}`
-);
-
-export function createAIPageRootBlockSpec(
-  framework: FrameworkProvider
-): ExtensionType[] {
-  return [...PageRootBlockSpec, aiPanelWidget, getAIPageRootWatcher(framework)];
-}
-
-function getAIEdgelessRootWatcher(framework: FrameworkProvider) {
-  class AIEdgelessRootWatcher extends BlockServiceWatcher {
-    static override readonly flavour = 'affine:page';
-
-    override mounted() {
-      super.mounted();
-      this.blockService.specSlots.widgetConnected.on(view => {
-        if (view.component instanceof AffineAIPanelWidget) {
-          view.component.style.width = '430px';
-          view.component.config = buildAIPanelConfig(view.component, framework);
-          setupSpaceAIEntry(view.component);
-        }
-
-        if (view.component instanceof EdgelessCopilotWidget) {
-          setupEdgelessCopilot(view.component);
-        }
-
-        if (view.component instanceof EdgelessElementToolbarWidget) {
-          setupEdgelessElementToolbarAIEntry(view.component);
-        }
-
-        if (view.component instanceof AffineFormatBarWidget) {
-          setupFormatBarAIEntry(view.component);
-        }
-
-        if (view.component instanceof AffineSlashMenuWidget) {
-          setupSlashMenuAIEntry(view.component);
-        }
-      });
-    }
-  }
-  return AIEdgelessRootWatcher;
-}
-
-export function createAIEdgelessRootBlockSpec(
-  framework: FrameworkProvider
-): ExtensionType[] {
-  return [
-    ...EdgelessRootBlockSpec,
-    CopilotTool,
-    aiPanelWidget,
-    edgelessCopilotWidget,
-    getAIEdgelessRootWatcher(framework),
-  ];
-}
-
-class AIParagraphBlockWatcher extends BlockServiceWatcher {
-  static override readonly flavour = 'affine:paragraph';
-
-  override mounted() {
-    super.mounted();
-    const service = this.blockService;
-    assertInstanceOf(service, ParagraphBlockService);
-    service.placeholderGenerator = model => {
-      if (model.type === 'text') {
-        return "Type '/' for commands, 'space' for AI";
-      }
-
-      const placeholders = {
-        h1: 'Heading 1',
-        h2: 'Heading 2',
-        h3: 'Heading 3',
-        h4: 'Heading 4',
-        h5: 'Heading 5',
-        h6: 'Heading 6',
-        quote: '',
-      };
-      return placeholders[model.type];
-    };
-  }
-}
-
-export const AIParagraphBlockSpec: ExtensionType[] = [
-  ...ParagraphBlockSpec,
-  AIParagraphBlockWatcher,
-];
-
-class AICodeBlockWatcher extends BlockServiceWatcher {
-  static override readonly flavour = 'affine:code';
-
-  override mounted() {
-    super.mounted();
-    const service = this.blockService;
-    service.specSlots.widgetConnected.on(view => {
-      if (view.component instanceof AffineCodeToolbarWidget) {
-        setupCodeToolbarAIEntry(view.component);
-      }
-    });
-  }
-}
-
-export const AICodeBlockSpec: ExtensionType[] = [
-  ...CodeBlockSpec,
-  AICodeBlockWatcher,
-];
-
-class AIImageBlockWatcher extends BlockServiceWatcher {
-  static override readonly flavour = 'affine:image';
-
-  override mounted() {
-    super.mounted();
-    this.blockService.specSlots.widgetConnected.on(view => {
-      if (view.component instanceof AffineImageToolbarWidget) {
-        setupImageToolbarAIEntry(view.component);
-      }
-    });
-  }
-}
-
-export const AIImageBlockSpec: ExtensionType[] = [
-  ...ImageBlockSpec,
-  AIImageBlockWatcher,
-];

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

@@ -0,0 +1,27 @@
+import { BlockServiceWatcher } from '@blocksuite/affine/block-std';
+import {
+  AffineCodeToolbarWidget,
+  CodeBlockSpec,
+} from '@blocksuite/affine/blocks';
+import type { ExtensionType } from '@blocksuite/affine/store';
+
+import { setupCodeToolbarAIEntry } from '../entries/code-toolbar/setup-code-toolbar';
+
+class AICodeBlockWatcher extends BlockServiceWatcher {
+  static override readonly flavour = 'affine:code';
+
+  override mounted() {
+    super.mounted();
+    const service = this.blockService;
+    service.specSlots.widgetConnected.on(view => {
+      if (view.component instanceof AffineCodeToolbarWidget) {
+        setupCodeToolbarAIEntry(view.component);
+      }
+    });
+  }
+}
+
+export const AICodeBlockSpec: ExtensionType[] = [
+  ...CodeBlockSpec,
+  AICodeBlockWatcher,
+];

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

@@ -0,0 +1,73 @@
+import { BlockServiceWatcher } from '@blocksuite/affine/block-std';
+import {
+  AffineFormatBarWidget,
+  AffineSlashMenuWidget,
+  EdgelessElementToolbarWidget,
+  EdgelessRootBlockSpec,
+} from '@blocksuite/affine/blocks';
+import type { ExtensionType } from '@blocksuite/affine/store';
+import type { FrameworkProvider } from '@toeverything/infra';
+
+import { buildAIPanelConfig } from '../ai-panel';
+import {
+  setupEdgelessCopilot,
+  setupEdgelessElementToolbarAIEntry,
+} from '../entries/edgeless/index';
+import { setupFormatBarAIEntry } from '../entries/format-bar/setup-format-bar';
+import { setupSlashMenuAIEntry } from '../entries/slash-menu/setup-slash-menu';
+import { setupSpaceAIEntry } from '../entries/space/setup-space';
+import { CopilotTool } from '../tool/copilot-tool';
+import {
+  AffineAIPanelWidget,
+  aiPanelWidget,
+} from '../widgets/ai-panel/ai-panel';
+import {
+  EdgelessCopilotWidget,
+  edgelessCopilotWidget,
+} from '../widgets/edgeless-copilot';
+
+export function createAIEdgelessRootBlockSpec(
+  framework: FrameworkProvider
+): ExtensionType[] {
+  return [
+    ...EdgelessRootBlockSpec,
+    CopilotTool,
+    aiPanelWidget,
+    edgelessCopilotWidget,
+    getAIEdgelessRootWatcher(framework),
+  ];
+}
+
+function getAIEdgelessRootWatcher(framework: FrameworkProvider) {
+  class AIEdgelessRootWatcher extends BlockServiceWatcher {
+    static override readonly flavour = 'affine:page';
+
+    override mounted() {
+      super.mounted();
+      this.blockService.specSlots.widgetConnected.on(view => {
+        if (view.component instanceof AffineAIPanelWidget) {
+          view.component.style.width = '430px';
+          view.component.config = buildAIPanelConfig(view.component, framework);
+          setupSpaceAIEntry(view.component);
+        }
+
+        if (view.component instanceof EdgelessCopilotWidget) {
+          setupEdgelessCopilot(view.component);
+        }
+
+        if (view.component instanceof EdgelessElementToolbarWidget) {
+          setupEdgelessElementToolbarAIEntry(view.component);
+        }
+
+        if (view.component instanceof AffineFormatBarWidget) {
+          setupFormatBarAIEntry(view.component);
+        }
+
+        if (view.component instanceof AffineSlashMenuWidget) {
+          setupSlashMenuAIEntry(view.component);
+        }
+      });
+    }
+  }
+  return AIEdgelessRootWatcher;
+}

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

@@ -0,0 +1,26 @@
+import { BlockServiceWatcher } from '@blocksuite/affine/block-std';
+import {
+  AffineImageToolbarWidget,
+  ImageBlockSpec,
+} from '@blocksuite/affine/blocks';
+import type { ExtensionType } from '@blocksuite/affine/store';
+
+import { setupImageToolbarAIEntry } from '../entries/image-toolbar/setup-image-toolbar';
+
+class AIImageBlockWatcher extends BlockServiceWatcher {
+  static override readonly flavour = 'affine:image';
+
+  override mounted() {
+    super.mounted();
+    this.blockService.specSlots.widgetConnected.on(view => {
+      if (view.component instanceof AffineImageToolbarWidget) {
+        setupImageToolbarAIEntry(view.component);
+      }
+    });
+  }
+}
+
+export const AIImageBlockSpec: ExtensionType[] = [
+  ...ImageBlockSpec,
+  AIImageBlockWatcher,
+];

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

@@ -0,0 +1,49 @@
+import { BlockServiceWatcher } from '@blocksuite/affine/block-std';
+import {
+  AffineFormatBarWidget,
+  AffineSlashMenuWidget,
+  PageRootBlockSpec,
+} from '@blocksuite/affine/blocks';
+import type { ExtensionType } from '@blocksuite/affine/store';
+import type { FrameworkProvider } from '@toeverything/infra';
+
+import { buildAIPanelConfig } from '../ai-panel';
+import { setupFormatBarAIEntry } from '../entries/format-bar/setup-format-bar';
+import { setupSlashMenuAIEntry } from '../entries/slash-menu/setup-slash-menu';
+import { setupSpaceAIEntry } from '../entries/space/setup-space';
+import {
+  AffineAIPanelWidget,
+  aiPanelWidget,
+} from '../widgets/ai-panel/ai-panel';
+
+function getAIPageRootWatcher(framework: FrameworkProvider) {
+  class AIPageRootWatcher extends BlockServiceWatcher {
+    static override readonly flavour = 'affine:page';
+
+    override mounted() {
+      super.mounted();
+      this.blockService.specSlots.widgetConnected.on(view => {
+        if (view.component instanceof AffineAIPanelWidget) {
+          view.component.style.width = '630px';
+          view.component.config = buildAIPanelConfig(view.component, framework);
+          setupSpaceAIEntry(view.component);
+        }
+
+        if (view.component instanceof AffineFormatBarWidget) {
+          setupFormatBarAIEntry(view.component);
+        }
+
+        if (view.component instanceof AffineSlashMenuWidget) {
+          setupSlashMenuAIEntry(view.component);
+        }
+      });
+    }
+  }
+  return AIPageRootWatcher;
+}
+
+export function createAIPageRootBlockSpec(
+  framework: FrameworkProvider
+): ExtensionType[] {
+  return [...PageRootBlockSpec, aiPanelWidget, getAIPageRootWatcher(framework)];
+}

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

@@ -0,0 +1,38 @@
+import { BlockServiceWatcher } from '@blocksuite/affine/block-std';
+import {
+  ParagraphBlockService,
+  ParagraphBlockSpec,
+} from '@blocksuite/affine/blocks';
+import { assertInstanceOf } from '@blocksuite/affine/global/utils';
+import type { ExtensionType } from '@blocksuite/affine/store';
+
+class AIParagraphBlockWatcher extends BlockServiceWatcher {
+  static override readonly flavour = 'affine:paragraph';
+
+  override mounted() {
+    super.mounted();
+    const service = this.blockService;
+    assertInstanceOf(service, ParagraphBlockService);
+    service.placeholderGenerator = model => {
+      if (model.type === 'text') {
+        return "Type '/' for commands, 'space' for AI";
+      }
+
+      const placeholders = {
+        h1: 'Heading 1',
+        h2: 'Heading 2',
+        h3: 'Heading 3',
+        h4: 'Heading 4',
+        h5: 'Heading 5',
+        h6: 'Heading 6',
+        quote: '',
+      };
+      return placeholders[model.type];
+    };
+  }
+}
+
+export const AIParagraphBlockSpec: ExtensionType[] = [
+  ...ParagraphBlockSpec,
+  AIParagraphBlockWatcher,
+];

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

@@ -0,0 +1,37 @@
+import {
+  CodeBlockSpec,
+  EdgelessRootBlockSpec,
+  ImageBlockSpec,
+  PageRootBlockSpec,
+  ParagraphBlockSpec,
+  type SpecBuilder,
+} from '@blocksuite/affine/blocks';
+import type { FrameworkProvider } from '@toeverything/infra';
+
+import { AIChatBlockSpec } from '../blocks';
+import { AICodeBlockSpec } from './ai-code';
+import { createAIEdgelessRootBlockSpec } from './ai-edgeless-root';
+import { AIImageBlockSpec } from './ai-image';
+import { createAIPageRootBlockSpec } from './ai-page-root';
+import { AIParagraphBlockSpec } from './ai-paragraph';
+
+export function enableAIExtension(
+  specBuilder: SpecBuilder,
+  framework: FrameworkProvider
+) {
+  specBuilder.replace(CodeBlockSpec, AICodeBlockSpec);
+  specBuilder.replace(ImageBlockSpec, AIImageBlockSpec);
+  specBuilder.replace(ParagraphBlockSpec, AIParagraphBlockSpec);
+
+  if (specBuilder.hasAll(EdgelessRootBlockSpec)) {
+    const aiEdgeless = createAIEdgelessRootBlockSpec(framework);
+    specBuilder.replace(EdgelessRootBlockSpec, aiEdgeless);
+  }
+
+  if (specBuilder.hasAll(PageRootBlockSpec)) {
+    const aiPage = createAIPageRootBlockSpec(framework);
+    specBuilder.replace(PageRootBlockSpec, aiPage);
+  }
+
+  specBuilder.extend(AIChatBlockSpec);
+}

packages/frontend/core/src/blocksuite/ai/extensions/index.ts

@@ -0,0 +1 @@
+export * from './enable-ai';

packages/frontend/core/src/blocksuite/ai/index.ts

@@ -1,10 +1,10 @@
 export * from './_common/config';
-export * from './actions/index';
-export * from './ai-spec';
-export { ChatPanel } from './chat-panel/index';
+export * from './actions';
+export { ChatPanel } from './chat-panel';
+export * from './entries';
 export * from './entries/edgeless/actions-config';
-export * from './entries/index';
-export * from './messages/index';
+export * from './extensions';
+export * from './messages';
 export { AIChatBlockPeekViewTemplate } from './peek-view/chat-block-peek-view';
 export * from './provider';
 export * from './utils/edgeless';

packages/frontend/core/src/blocksuite/ai/widgets/ai-panel/ai-panel.ts

@@ -1,4 +1,7 @@
-import { WidgetComponent } from '@blocksuite/affine/block-std';
+import {
+  WidgetComponent,
+  WidgetViewExtension,
+} from '@blocksuite/affine/block-std';
 import { GfxControllerIdentifier } from '@blocksuite/affine/block-std/gfx';
 import {
   AFFINE_FORMAT_BAR_WIDGET,
@@ -25,6 +28,7 @@ import {
 import { css, html, nothing, type PropertyValues } from 'lit';
 import { property, query } from 'lit/decorators.js';
 import { choose } from 'lit/directives/choose.js';
+import { literal, unsafeStatic } from 'lit/static-html.js';
 
 import type { AIError } from '../../components/ai-item/types.js';
 import type { AIPanelGenerating } from './components/index.js';
@@ -543,3 +547,9 @@ export class AffineAIPanelWidget extends WidgetComponent {
   @property()
   accessor state: AffineAIPanelState = 'hidden';
 }
+
+export const aiPanelWidget = WidgetViewExtension(
+  'affine:page',
+  AFFINE_AI_PANEL_WIDGET,
+  literal`${unsafeStatic(AFFINE_AI_PANEL_WIDGET)}`
+);

packages/frontend/core/src/blocksuite/ai/widgets/edgeless-copilot/index.ts

@@ -1,4 +1,7 @@
-import { WidgetComponent } from '@blocksuite/affine/block-std';
+import {
+  WidgetComponent,
+  WidgetViewExtension,
+} from '@blocksuite/affine/block-std';
 import { GfxControllerIdentifier } from '@blocksuite/affine/block-std/gfx';
 import type { RootBlockModel } from '@blocksuite/affine/blocks';
 import {
@@ -22,6 +25,7 @@ import { effect } from '@preact/signals-core';
 import { css, html, nothing } from 'lit';
 import { query, state } from 'lit/decorators.js';
 import { styleMap } from 'lit/directives/style-map.js';
+import { literal, unsafeStatic } from 'lit/static-html.js';
 
 import type { AIItemGroupConfig } from '../../components/ai-item/types.js';
 import {
@@ -290,6 +294,12 @@ export class EdgelessCopilotWidget extends WidgetComponent<RootBlockModel> {
   accessor selectionElem!: HTMLDivElement;
 }
 
+export const edgelessCopilotWidget = WidgetViewExtension(
+  'affine:page',
+  AFFINE_EDGELESS_COPILOT_WIDGET,
+  literal`${unsafeStatic(AFFINE_EDGELESS_COPILOT_WIDGET)}`
+);
+
 declare global {
   interface HTMLElementTagNameMap {
     [AFFINE_EDGELESS_COPILOT_WIDGET]: EdgelessCopilotWidget;

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

@@ -1,60 +0,0 @@
-import {
-  AICodeBlockSpec,
-  AIImageBlockSpec,
-  AIParagraphBlockSpec,
-} from '@affine/core/blocksuite/ai';
-import { AIChatBlockSpec } from '@affine/core/blocksuite/ai/blocks';
-import {
-  AdapterFactoryExtensions,
-  AttachmentBlockSpec,
-  BookmarkBlockSpec,
-  CodeBlockSpec,
-  DatabaseBlockSpec,
-  DataViewBlockSpec,
-  DefaultOpenDocExtension,
-  DividerBlockSpec,
-  EditPropsStore,
-  EmbedExtensions,
-  FontLoaderService,
-  ImageBlockSpec,
-  LatexBlockSpec,
-  ListBlockSpec,
-  ParagraphBlockSpec,
-  RefNodeSlotsExtension,
-  RichTextExtensions,
-  TableBlockSpec,
-} from '@blocksuite/affine/blocks';
-import type { ExtensionType } from '@blocksuite/affine/store';
-
-const CommonBlockSpecs: ExtensionType[] = [
-  RefNodeSlotsExtension,
-  EditPropsStore,
-  RichTextExtensions,
-  LatexBlockSpec,
-  ListBlockSpec,
-  DatabaseBlockSpec,
-  TableBlockSpec,
-  DataViewBlockSpec,
-  DividerBlockSpec,
-  EmbedExtensions,
-  BookmarkBlockSpec,
-  AttachmentBlockSpec,
-  AdapterFactoryExtensions,
-  FontLoaderService,
-  DefaultOpenDocExtension,
-].flat();
-
-export const DefaultBlockSpecs: ExtensionType[] = [
-  CodeBlockSpec,
-  ImageBlockSpec,
-  ParagraphBlockSpec,
-  ...CommonBlockSpecs,
-].flat();
-
-export const AIBlockSpecs: ExtensionType[] = [
-  AICodeBlockSpec,
-  AIImageBlockSpec,
-  AIParagraphBlockSpec,
-  AIChatBlockSpec,
-  ...CommonBlockSpecs,
-].flat();

packages/frontend/core/src/blocksuite/block-suite-editor/specs/custom/root-block.ts

@@ -1,9 +1,3 @@
-import {
-  AICodeBlockSpec,
-  AIImageBlockSpec,
-  AIParagraphBlockSpec,
-} from '@affine/core/blocksuite/ai';
-import { AIChatBlockSpec } from '@affine/core/blocksuite/ai/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';
@@ -19,14 +13,11 @@ import type {
   ThemeExtension,
 } from '@blocksuite/affine/blocks';
 import {
-  CodeBlockSpec,
   ColorScheme,
   createSignalFromObservable,
   DatabaseConfigExtension,
   DocDisplayMetaProvider,
   EditorSettingExtension,
-  ImageBlockSpec,
-  ParagraphBlockSpec,
   referenceToNode,
   RootBlockConfigExtension,
   SpecProvider,
@@ -267,10 +258,3 @@ export function enableAffineExtension(
     ].flat()
   );
 }
-
-export function enableAIExtension(specBuilder: SpecBuilder): void {
-  specBuilder.replace(CodeBlockSpec, AICodeBlockSpec);
-  specBuilder.replace(ImageBlockSpec, AIImageBlockSpec);
-  specBuilder.replace(ParagraphBlockSpec, AIParagraphBlockSpec);
-  specBuilder.extend(AIChatBlockSpec);
-}

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

@@ -1,16 +1,12 @@
-import { createAIEdgelessRootBlockSpec } from '@affine/core/blocksuite/ai';
+import { enableAIExtension } from '@affine/core/blocksuite/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 { SpecBuilder, TemplateManager } from '@blocksuite/affine/blocks';
-import {
-  EdgelessRootBlockSpec,
-  EdgelessTemplatePanel,
-  SpecProvider,
-} from '@blocksuite/affine/blocks';
+import { EdgelessTemplatePanel, SpecProvider } from '@blocksuite/affine/blocks';
 import { type FrameworkProvider } from '@toeverything/infra';
 
-import { enableAffineExtension, enableAIExtension } from './custom/root-block';
+import { enableAffineExtension } from './custom/root-block';
 
 export function createEdgelessModeSpecs(
   framework: FrameworkProvider
@@ -20,11 +16,7 @@ export function createEdgelessModeSpecs(
   const edgelessSpec = SpecProvider._.getSpec('edgeless');
   enableAffineExtension(framework, edgelessSpec);
   if (enableAI) {
-    enableAIExtension(edgelessSpec);
-    edgelessSpec.replace(
-      EdgelessRootBlockSpec,
-      createAIEdgelessRootBlockSpec(framework)
-    );
+    enableAIExtension(edgelessSpec, framework);
   }
 
   return edgelessSpec;

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

@@ -1,13 +1,9 @@
-import { createAIPageRootBlockSpec } from '@affine/core/blocksuite/ai';
+import { enableAIExtension } from '@affine/core/blocksuite/ai';
 import { FeatureFlagService } from '@affine/core/modules/feature-flag';
-import {
-  PageRootBlockSpec,
-  type SpecBuilder,
-  SpecProvider,
-} from '@blocksuite/affine/blocks';
+import { type SpecBuilder, SpecProvider } from '@blocksuite/affine/blocks';
 import { type FrameworkProvider } from '@toeverything/infra';
 
-import { enableAffineExtension, enableAIExtension } from './custom/root-block';
+import { enableAffineExtension } from './custom/root-block';
 
 export function createPageModeSpecs(framework: FrameworkProvider): SpecBuilder {
   const featureFlagService = framework.get(FeatureFlagService);
@@ -16,8 +12,7 @@ export function createPageModeSpecs(framework: FrameworkProvider): SpecBuilder {
   const pageSpec = provider.getSpec('page');
   enableAffineExtension(framework, pageSpec);
   if (enableAI) {
-    enableAIExtension(pageSpec);
-    pageSpec.replace(PageRootBlockSpec, createAIPageRootBlockSpec(framework));
+    enableAIExtension(pageSpec, framework);
   }
   return pageSpec;
 }