feat(core): impl AI switch (#8018)

close PD-1658

https://github.com/user-attachments/assets/2f3d1d26-1879-4d95-b80c-7c0965cefbd0

packages/frontend/core/src/components/affine/ai-onboarding/index.tsx

@@ -1,3 +1,5 @@
+import { EditorSettingService } from '@affine/core/modules/editor-settting';
+import { useLiveData, useService } from '@toeverything/infra';
 import { Suspense, useCallback, useEffect, useState } from 'react';
 
 import { AIOnboardingEdgeless } from './edgeless.dialog';
@@ -28,19 +30,29 @@ const useDismiss = (key: AIOnboardingType) => {
 export const WorkspaceAIOnboarding = () => {
   const [dismissGeneral] = useDismiss(AIOnboardingType.GENERAL);
   const [dismissLocal] = useDismiss(AIOnboardingType.LOCAL);
+  const editorSettingService = useService(EditorSettingService);
+  const enableAI = useLiveData(
+    editorSettingService.editorSetting.settings$.map(s => s.enableAI)
+  );
 
   return (
     <Suspense>
-      {dismissGeneral ? null : <AIOnboardingGeneral />}
-      {dismissLocal ? null : <AIOnboardingLocal />}
+      {!enableAI || dismissGeneral ? null : <AIOnboardingGeneral />}
+      {!enableAI || dismissLocal ? null : <AIOnboardingLocal />}
     </Suspense>
   );
 };
 
 export const PageAIOnboarding = () => {
   const [dismissEdgeless] = useDismiss(AIOnboardingType.EDGELESS);
+  const editorSettingService = useService(EditorSettingService);
+  const enableAI = useLiveData(
+    editorSettingService.editorSetting.settings$.map(s => s.enableAI)
+  );
 
   return (
-    <Suspense>{dismissEdgeless ? null : <AIOnboardingEdgeless />}</Suspense>
+    <Suspense>
+      {!enableAI || dismissEdgeless ? null : <AIOnboardingEdgeless />}
+    </Suspense>
   );
 };

packages/frontend/core/src/components/affine/setting-modal/general-setting/editor/general.tsx

@@ -8,6 +8,7 @@ import {
   type RadioItem,
   Scrollable,
   Switch,
+  useConfirmModal,
 } from '@affine/component';
 import {
   SettingRow,
@@ -393,16 +394,51 @@ export const SpellCheckSettings = () => {
   );
 };
 
+const AISettings = () => {
+  const t = useI18n();
+  const { openConfirmModal } = useConfirmModal();
+  const { editorSettingService } = useServices({ EditorSettingService });
+
+  const settings = useLiveData(editorSettingService.editorSetting.settings$);
+
+  const onAIChange = useCallback(
+    (checked: boolean) => {
+      editorSettingService.editorSetting.set('enableAI', checked);
+    },
+    [editorSettingService]
+  );
+  const onToggleAI = useCallback(
+    (checked: boolean) => {
+      openConfirmModal({
+        title: checked ? 'Enable AI?' : 'Disable AI?',
+        description: `Are you sure you want to ${checked ? 'enable' : 'disable'} AI?`,
+        confirmText: checked ? 'Enable' : 'Disable',
+        cancelText: 'Cancel',
+        onConfirm: () => onAIChange(checked),
+        confirmButtonOptions: {
+          variant: checked ? 'primary' : 'error',
+        },
+      });
+    },
+    [openConfirmModal, onAIChange]
+  );
+
+  return (
+    <SettingRow
+      name={t['com.affine.settings.editorSettings.general.ai.title']()}
+      desc={t['com.affine.settings.editorSettings.general.ai.description']()}
+    >
+      <Switch checked={settings.enableAI} onChange={onToggleAI} />
+    </SettingRow>
+  );
+};
+
 export const General = () => {
   const t = useI18n();
+
   return (
     <SettingWrapper title={t['com.affine.settings.editorSettings.general']()}>
-      <SettingRow
-        name={t['com.affine.settings.editorSettings.general.ai.title']()}
-        desc={t['com.affine.settings.editorSettings.general.ai.description']()}
-      >
-        <Switch />
-      </SettingRow>
+      <AISettings />
       <FontFamilySettings />
       <CustomFontFamilySettings />
       <NewDocDefaultModeSettings />

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

@@ -15,6 +15,7 @@ import {
   useFramework,
   useLiveData,
   useService,
+  useServices,
 } from '@toeverything/infra';
 import React, {
   forwardRef,
@@ -65,9 +66,13 @@ interface BlocksuiteEditorProps {
 
 const usePatchSpecs = (page: Doc, shared: boolean, mode: DocMode) => {
   const [reactToLit, portals] = useLitPortalFactory();
-  const peekViewService = useService(PeekViewService);
-  const docService = useService(DocService);
-  const docsService = useService(DocsService);
+  const { peekViewService, docService, docsService, editorSettingService } =
+    useServices({
+      PeekViewService,
+      DocService,
+      DocsService,
+      EditorSettingService,
+    });
   const framework = useFramework();
   const referenceRenderer: ReferenceReactRenderer = useMemo(() => {
     return function customReference(reference) {
@@ -89,10 +94,12 @@ const usePatchSpecs = (page: Doc, shared: boolean, mode: DocMode) => {
   }, [mode, page.collection]);
 
   const specs = useMemo(() => {
+    const enableAI =
+      editorSettingService.editorSetting.settings$.value.enableAI;
     return mode === 'edgeless'
-      ? createEdgelessModeSpecs(framework)
-      : createPageModeSpecs(framework);
-  }, [mode, framework]);
+      ? createEdgelessModeSpecs(framework, enableAI)
+      : createPageModeSpecs(framework, enableAI);
+  }, [editorSettingService, mode, framework]);
 
   const confirmModal = useConfirmModal();
   const patchedSpecs = useMemo(() => {

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

@@ -6,6 +6,7 @@ import {
 import type { ExtensionType } from '@blocksuite/block-std';
 import {
   BookmarkBlockSpec,
+  CodeBlockSpec,
   DatabaseBlockSpec,
   DataViewBlockSpec,
   DividerBlockSpec,
@@ -16,12 +17,15 @@ import {
   EmbedLoomBlockSpec,
   EmbedSyncedDocBlockSpec,
   EmbedYoutubeBlockSpec,
+  ImageBlockSpec,
   ListBlockSpec,
+  ParagraphBlockSpec,
 } from '@blocksuite/blocks';
+import { AIChatBlockSpec } from '@blocksuite/presets';
 
 import { CustomAttachmentBlockSpec } from './custom/attachment-block';
 
-export const CommonBlockSpecs: ExtensionType[] = [
+const CommonBlockSpecs: ExtensionType[] = [
   ListBlockSpec,
   DatabaseBlockSpec,
   DataViewBlockSpec,
@@ -36,7 +40,19 @@ export const CommonBlockSpecs: ExtensionType[] = [
   EmbedLinkedDocBlockSpec,
   // special
   CustomAttachmentBlockSpec,
+].flat();
+
+export const DefaultBlockSpecs: ExtensionType[] = [
+  CodeBlockSpec,
+  ImageBlockSpec,
+  ParagraphBlockSpec,
+  ...CommonBlockSpecs,
+].flat();
+
+export const AIBlockSpecs: ExtensionType[] = [
   AICodeBlockSpec,
   AIImageBlockSpec,
   AIParagraphBlockSpec,
+  AIChatBlockSpec,
+  ...CommonBlockSpecs,
 ].flat();

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

@@ -14,7 +14,9 @@ import {
 import type { RootService, TelemetryEventMap } from '@blocksuite/blocks';
 import {
   AffineCanvasTextFonts,
+  EdgelessRootBlockSpec,
   EdgelessRootService,
+  PageRootBlockSpec,
   PageRootService,
 } from '@blocksuite/blocks';
 import { type FrameworkProvider } from '@toeverything/infra';
@@ -55,11 +57,12 @@ function withAffineRootService(Service: typeof PageRootService) {
 }
 
 export function createPageRootBlockSpec(
-  framework: FrameworkProvider
+  framework: FrameworkProvider,
+  enableAI: boolean
 ): ExtensionType[] {
   const editorSettingService = framework.get(EditorSettingService);
   return [
-    ...AIPageRootBlockSpec,
+    ...(enableAI ? AIPageRootBlockSpec : PageRootBlockSpec),
     {
       setup: di => {
         di.override(
@@ -79,11 +82,12 @@ export function createPageRootBlockSpec(
 }
 
 export function createEdgelessRootBlockSpec(
-  framework: FrameworkProvider
+  framework: FrameworkProvider,
+  enableAI: boolean
 ): ExtensionType[] {
   const editorSettingService = framework.get(EditorSettingService);
   return [
-    ...AIEdgelessRootBlockSpec,
+    ...(enableAI ? AIEdgelessRootBlockSpec : EdgelessRootBlockSpec),
     {
       setup: di => {
         di.override(

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

@@ -6,24 +6,23 @@ import {
   EdgelessTextBlockSpec,
   FrameBlockSpec,
 } from '@blocksuite/blocks';
-import { AIChatBlockSpec } from '@blocksuite/presets';
 import type { FrameworkProvider } from '@toeverything/infra';
 
-import { CommonBlockSpecs } from './common';
+import { AIBlockSpecs, DefaultBlockSpecs } from './common';
 import { createEdgelessRootBlockSpec } from './custom/root-block';
 
 export function createEdgelessModeSpecs(
-  framework: FrameworkProvider
+  framework: FrameworkProvider,
+  enableAI: boolean
 ): ExtensionType[] {
   return [
-    ...CommonBlockSpecs,
+    ...(enableAI ? AIBlockSpecs : DefaultBlockSpecs),
     EdgelessSurfaceBlockSpec,
     EdgelessSurfaceRefBlockSpec,
     FrameBlockSpec,
     EdgelessTextBlockSpec,
     EdgelessNoteBlockSpec,
-    AIChatBlockSpec,
     // special
-    createEdgelessRootBlockSpec(framework),
+    createEdgelessRootBlockSpec(framework, enableAI),
   ].flat();
 }

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

@@ -4,22 +4,21 @@ import {
   PageSurfaceBlockSpec,
   PageSurfaceRefBlockSpec,
 } from '@blocksuite/blocks';
-import { AIChatBlockSpec } from '@blocksuite/presets';
 import { type FrameworkProvider } from '@toeverything/infra';
 
-import { CommonBlockSpecs } from './common';
+import { AIBlockSpecs, DefaultBlockSpecs } from './common';
 import { createPageRootBlockSpec } from './custom/root-block';
 
 export function createPageModeSpecs(
-  framework: FrameworkProvider
+  framework: FrameworkProvider,
+  enableAI: boolean
 ): ExtensionType[] {
   return [
-    ...CommonBlockSpecs,
+    ...(enableAI ? AIBlockSpecs : DefaultBlockSpecs),
     PageSurfaceBlockSpec,
     PageSurfaceRefBlockSpec,
     NoteBlockSpec,
-    AIChatBlockSpec,
     // special
-    createPageRootBlockSpec(framework),
+    createPageRootBlockSpec(framework, enableAI),
   ].flat();
 }