refactor(editor): extract ai extension builder (#12240)

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

@@ -165,7 +165,8 @@ const usePreviewExtensions = () => {
   const extensions = useMemo(() => {
     const manager = getViewManager()
       .config.init()
-      .common(framework, enableAI)
+      .common(framework)
+      .ai(enableAI, framework)
       .theme(framework)
       .database(framework)
       .linkedDoc(framework)

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

@@ -80,7 +80,8 @@ const usePatchSpecs = (mode: DocMode) => {
   const patchedSpecs = useMemo(() => {
     const manager = getViewManager()
       .config.init()
-      .common(framework, enableAI)
+      .common(framework)
+      .ai(enableAI, framework)
       .theme(framework)
       .editorConfig(framework)
       .editorView({
@@ -99,7 +100,8 @@ const usePatchSpecs = (mode: DocMode) => {
       .database(framework)
       .linkedDoc(framework)
       .paragraph(enableAI)
-      .mobile(framework).value;
+      .mobile(framework)
+      .electron(framework).value;
 
     if (BUILD_CONFIG.isMobileEdition) {
       if (mode === 'page') {

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

@@ -0,0 +1,79 @@
+import { toolbarAIEntryConfig } from '@affine/core/blocksuite/ai';
+import { AIChatBlockSpec } from '@affine/core/blocksuite/ai/blocks';
+import { AITranscriptionBlockSpec } from '@affine/core/blocksuite/ai/blocks/ai-chat-block/ai-transcription-block';
+import { edgelessToolbarAIEntryConfig } from '@affine/core/blocksuite/ai/entries/edgeless';
+import { imageToolbarAIEntryConfig } from '@affine/core/blocksuite/ai/entries/image-toolbar/setup-image-toolbar';
+import { AICodeBlockWatcher } from '@affine/core/blocksuite/ai/extensions/ai-code';
+import { getAIEdgelessRootWatcher } from '@affine/core/blocksuite/ai/extensions/ai-edgeless-root';
+import { getAIPageRootWatcher } from '@affine/core/blocksuite/ai/extensions/ai-page-root';
+import { AiSlashMenuConfigExtension } from '@affine/core/blocksuite/ai/extensions/ai-slash-menu';
+import { CopilotTool } from '@affine/core/blocksuite/ai/tool/copilot-tool';
+import { aiPanelWidget } from '@affine/core/blocksuite/ai/widgets/ai-panel/ai-panel';
+import { edgelessCopilotWidget } from '@affine/core/blocksuite/ai/widgets/edgeless-copilot';
+import {
+  type ViewExtensionContext,
+  ViewExtensionProvider,
+} from '@blocksuite/affine/ext-loader';
+import { ToolbarModuleExtension } from '@blocksuite/affine/shared/services';
+import { BlockFlavourIdentifier } from '@blocksuite/affine/std';
+import { FrameworkProvider } from '@toeverything/infra';
+import { z } from 'zod';
+
+import { EdgelessClipboardAIChatConfig } from './edgeless-clipboard';
+
+const optionsSchema = z.object({
+  enable: z.boolean().optional(),
+  framework: z.instanceof(FrameworkProvider).optional(),
+});
+
+type AIViewOptions = z.infer<typeof optionsSchema>;
+
+export class AIViewExtension extends ViewExtensionProvider<AIViewOptions> {
+  override name = 'affine-ai-view-extension';
+
+  override schema = optionsSchema;
+
+  override setup(context: ViewExtensionContext, options?: AIViewOptions) {
+    super.setup(context, options);
+    if (!options?.enable) return;
+    const framework = options.framework;
+    if (!framework) return;
+
+    context
+      .register(AIChatBlockSpec)
+      .register(AITranscriptionBlockSpec)
+      .register(EdgelessClipboardAIChatConfig)
+      .register(AICodeBlockWatcher)
+      .register(
+        ToolbarModuleExtension({
+          id: BlockFlavourIdentifier('custom:affine:image'),
+          config: imageToolbarAIEntryConfig(),
+        })
+      );
+    if (context.scope === 'edgeless' || context.scope === 'page') {
+      context.register([
+        aiPanelWidget,
+        AiSlashMenuConfigExtension(),
+        ToolbarModuleExtension({
+          id: BlockFlavourIdentifier('custom:affine:note'),
+          config: toolbarAIEntryConfig(),
+        }),
+      ]);
+    }
+    if (context.scope === 'edgeless') {
+      context.register([
+        CopilotTool,
+        edgelessCopilotWidget,
+        getAIEdgelessRootWatcher(framework),
+        // In note
+        ToolbarModuleExtension({
+          id: BlockFlavourIdentifier('custom:affine:surface:*'),
+          config: edgelessToolbarAIEntryConfig(),
+        }),
+      ]);
+    }
+    if (context.scope === 'page') {
+      context.register(getAIPageRootWatcher(framework));
+    }
+  }
+}

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

@@ -0,0 +1,33 @@
+import {
+  type ViewExtensionContext,
+  ViewExtensionProvider,
+} from '@blocksuite/affine/ext-loader';
+import { FrameworkProvider } from '@toeverything/infra';
+import { z } from 'zod';
+
+import { patchForClipboardInElectron } from './electron-clipboard';
+
+const optionsSchema = z.object({
+  framework: z.instanceof(FrameworkProvider).optional(),
+});
+
+type ElectronViewExtensionOptions = z.infer<typeof optionsSchema>;
+
+export class ElectronViewExtension extends ViewExtensionProvider<ElectronViewExtensionOptions> {
+  override name = 'electron-view-extensions';
+
+  override schema = optionsSchema;
+
+  override setup(
+    context: ViewExtensionContext,
+    options?: ElectronViewExtensionOptions
+  ) {
+    super.setup(context, options);
+    if (!BUILD_CONFIG.isElectron) return;
+
+    const framework = options?.framework;
+    if (!framework) return;
+
+    context.register(patchForClipboardInElectron(framework));
+  }
+}

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

@@ -1,17 +1,4 @@
-import { toolbarAIEntryConfig } from '@affine/core/blocksuite/ai';
-import { AIChatBlockSpec } from '@affine/core/blocksuite/ai/blocks';
-import { AITranscriptionBlockSpec } from '@affine/core/blocksuite/ai/blocks/ai-chat-block/ai-transcription-block';
-import { edgelessToolbarAIEntryConfig } from '@affine/core/blocksuite/ai/entries/edgeless';
-import { imageToolbarAIEntryConfig } from '@affine/core/blocksuite/ai/entries/image-toolbar/setup-image-toolbar';
-import { AICodeBlockWatcher } from '@affine/core/blocksuite/ai/extensions/ai-code';
-import { getAIEdgelessRootWatcher } from '@affine/core/blocksuite/ai/extensions/ai-edgeless-root';
-import { getAIPageRootWatcher } from '@affine/core/blocksuite/ai/extensions/ai-page-root';
-import { AiSlashMenuConfigExtension } from '@affine/core/blocksuite/ai/extensions/ai-slash-menu';
-import { CopilotTool } from '@affine/core/blocksuite/ai/tool/copilot-tool';
-import { aiPanelWidget } from '@affine/core/blocksuite/ai/widgets/ai-panel/ai-panel';
-import { edgelessCopilotWidget } from '@affine/core/blocksuite/ai/widgets/edgeless-copilot';
 import { buildDocDisplayMetaExtension } from '@affine/core/blocksuite/extensions/display-meta';
-import { EdgelessClipboardAIChatConfig } from '@affine/core/blocksuite/extensions/edgeless-clipboard';
 import { patchFileSizeLimitExtension } from '@affine/core/blocksuite/extensions/file-size-limit';
 import { getFontConfigExtension } from '@affine/core/blocksuite/extensions/font-config';
 import { patchPeekViewService } from '@affine/core/blocksuite/extensions/peek-view-service';
@@ -21,13 +8,10 @@ import {
   type ViewExtensionContext,
   ViewExtensionProvider,
 } from '@blocksuite/affine/ext-loader';
-import { ToolbarModuleExtension } from '@blocksuite/affine/shared/services';
-import { BlockFlavourIdentifier } from '@blocksuite/affine/std';
 import { FrameworkProvider } from '@toeverything/infra';
 import { z } from 'zod';
 
 const optionsSchema = z.object({
-  enableAI: z.boolean().optional(),
   framework: z.instanceof(FrameworkProvider).optional(),
 });
 
@@ -38,54 +22,12 @@ export class AffineCommonViewExtension extends ViewExtensionProvider<
 
   override schema = optionsSchema;
 
-  private _setupAI(
-    context: ViewExtensionContext,
-    framework: FrameworkProvider
-  ) {
-    context
-      .register(AIChatBlockSpec)
-      .register(AITranscriptionBlockSpec)
-      .register(EdgelessClipboardAIChatConfig)
-      .register(AICodeBlockWatcher)
-      .register(
-        ToolbarModuleExtension({
-          id: BlockFlavourIdentifier('custom:affine:image'),
-          config: imageToolbarAIEntryConfig(),
-        })
-      );
-    if (context.scope === 'edgeless' || context.scope === 'page') {
-      context.register([
-        aiPanelWidget,
-        AiSlashMenuConfigExtension(),
-        ToolbarModuleExtension({
-          id: BlockFlavourIdentifier('custom:affine:note'),
-          config: toolbarAIEntryConfig(),
-        }),
-      ]);
-    }
-    if (context.scope === 'edgeless') {
-      context.register([
-        CopilotTool,
-        edgelessCopilotWidget,
-        getAIEdgelessRootWatcher(framework),
-        // In note
-        ToolbarModuleExtension({
-          id: BlockFlavourIdentifier('custom:affine:surface:*'),
-          config: edgelessToolbarAIEntryConfig(),
-        }),
-      ]);
-    }
-    if (context.scope === 'page') {
-      context.register(getAIPageRootWatcher(framework));
-    }
-  }
-
   override setup(
     context: ViewExtensionContext,
     options?: z.infer<typeof optionsSchema>
   ) {
     super.setup(context, options);
-    const { framework, enableAI } = options || {};
+    const { framework } = options || {};
     if (framework) {
       context.register(patchPeekViewService(framework.get(PeekViewService)));
       context.register([
@@ -96,9 +38,6 @@ export class AffineCommonViewExtension extends ViewExtensionProvider<
       if (context.scope === 'edgeless' || context.scope === 'page') {
         context.register(patchFileSizeLimitExtension(framework));
       }
-      if (enableAI) {
-        this._setupAI(context, framework);
-      }
     }
   }
 }

packages/frontend/core/src/blocksuite/manager/editor-view.tsx

@@ -3,7 +3,6 @@ import { patchForAudioEmbedView } from '@affine/core/blocksuite/extensions/audio
 import { patchDatabaseBlockConfigService } from '@affine/core/blocksuite/extensions/database-block-config-service';
 import { patchDocModeService } from '@affine/core/blocksuite/extensions/doc-mode-service';
 import { patchDocUrlExtensions } from '@affine/core/blocksuite/extensions/doc-url';
-import { patchForClipboardInElectron } from '@affine/core/blocksuite/extensions/electron-clipboard';
 import { patchNotificationService } from '@affine/core/blocksuite/extensions/notification-service';
 import { patchOpenDocExtension } from '@affine/core/blocksuite/extensions/open-doc';
 import { patchQuickSearchService } from '@affine/core/blocksuite/extensions/quick-search-service';
@@ -98,7 +97,6 @@ export class AffineEditorViewExtension extends ViewExtensionProvider<AffineEdito
       reactToLit,
       confirmModal,
     } = options;
-    const isElectron = BUILD_CONFIG.isElectron;
 
     const docService = framework.get(DocService);
     const docsService = framework.get(DocsService);
@@ -106,21 +104,19 @@ export class AffineEditorViewExtension extends ViewExtensionProvider<AffineEdito
 
     const referenceRenderer = this._getCustomReferenceRenderer(framework);
 
-    context.register([
+    context
-      patchReferenceRenderer(reactToLit, referenceRenderer),
+      .register([
-      patchNotificationService(confirmModal),
+        patchReferenceRenderer(reactToLit, referenceRenderer),
-      patchOpenDocExtension(),
+        patchNotificationService(confirmModal),
-      patchSideBarService(framework),
+        patchOpenDocExtension(),
-      patchDocModeService(docService, docsService, editorService),
+        patchSideBarService(framework),
-    ]);
+        patchDocModeService(docService, docsService, editorService),
-    context.register(patchDocUrlExtensions(framework));
+      ])
-    context.register(patchQuickSearchService(framework));
+      .register(patchDocUrlExtensions(framework))
-    context.register([
+      .register(patchQuickSearchService(framework))
-      patchDatabaseBlockConfigService(),
+      .register([
-      patchForAudioEmbedView(reactToLit),
+        patchDatabaseBlockConfigService(),
-    ]);
+        patchForAudioEmbedView(reactToLit),
-    if (isElectron) {
+      ]);
-      context.register(patchForClipboardInElectron(framework));
-    }
   }
 }

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

@@ -1,4 +1,5 @@
 import type { ReactToLit } from '@affine/component';
+import { AIViewExtension } from '@affine/core/blocksuite/extensions/ai';
 import { CloudViewExtension } from '@affine/core/blocksuite/extensions/cloud';
 import {
   EdgelessBlockHeaderConfigViewExtension,
@@ -7,6 +8,7 @@ import {
 import { AffineEditorConfigViewExtension } from '@affine/core/blocksuite/extensions/editor-config';
 import { createDatabaseOptionsConfig } from '@affine/core/blocksuite/extensions/editor-config/database';
 import { createLinkedWidgetConfig } from '@affine/core/blocksuite/extensions/editor-config/linked';
+import { ElectronViewExtension } from '@affine/core/blocksuite/extensions/electron';
 import { MobileViewExtension } from '@affine/core/blocksuite/extensions/mobile';
 import { PdfViewExtension } from '@affine/core/blocksuite/extensions/pdf';
 import { AffineThemeViewExtension } from '@affine/core/blocksuite/extensions/theme';
@@ -40,6 +42,8 @@ type Configure = {
   turboRenderer: (enableTurboRenderer?: boolean) => Configure;
   pdf: (enablePDFEmbedPreview?: boolean, reactToLit?: ReactToLit) => Configure;
   mobile: (framework?: FrameworkProvider) => Configure;
+  ai: (enable?: boolean, framework?: FrameworkProvider) => Configure;
+  electron: (framework?: FrameworkProvider) => Configure;
 
   value: ViewExtensionManager;
 };
@@ -69,6 +73,8 @@ class ViewProvider {
       CloudViewExtension,
       PdfViewExtension,
       MobileViewExtension,
+      AIViewExtension,
+      ElectronViewExtension,
     ]);
   }
 
@@ -91,6 +97,8 @@ class ViewProvider {
       turboRenderer: this._configureTurboRenderer,
       pdf: this._configurePdf,
       mobile: this._configureMobile,
+      ai: this._configureAI,
+      electron: this._configureElectron,
       value: this._manager,
     };
   }
@@ -108,18 +116,16 @@ class ViewProvider {
       .cloud()
       .turboRenderer()
       .pdf()
-      .mobile();
+      .mobile()
+      .ai()
+      .electron();
 
     return this.config;
   };
 
-  private readonly _configureCommon = (
+  private readonly _configureCommon = (framework?: FrameworkProvider) => {
-    framework?: FrameworkProvider,
-    enableAI?: boolean
-  ) => {
     this._manager.configure(AffineCommonViewExtension, {
       framework,
-      enableAI,
     });
     return this.config;
   };
@@ -237,6 +243,19 @@ class ViewProvider {
     this._manager.configure(MobileViewExtension, { framework });
     return this.config;
   };
+
+  private readonly _configureAI = (
+    enable?: boolean,
+    framework?: FrameworkProvider
+  ) => {
+    this._manager.configure(AIViewExtension, { framework, enable });
+    return this.config;
+  };
+
+  private readonly _configureElectron = (framework?: FrameworkProvider) => {
+    this._manager.configure(ElectronViewExtension, { framework });
+    return this.config;
+  };
 }
 
 export function getViewManager() {