From 82916e8264a59a64e5676f88825f0193808e79bd Mon Sep 17 00:00:00 2001
From: JimmFly <yangjinfei001@gmail.com>
Date: Wed, 16 Oct 2024 13:03:56 +0800
Subject: [PATCH] feat(core): make AI functions follow server configuration
 (#8374)

---
 .../components/affine/ai-onboarding/index.tsx | 21 ++++++++++++++++---
 .../general-setting/editor/general.tsx        | 14 +++++++++++--
 .../block-suite-editor/lit-adaper.tsx         | 20 ++++++++++++++----
 3 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/packages/frontend/core/src/components/affine/ai-onboarding/index.tsx b/packages/frontend/core/src/components/affine/ai-onboarding/index.tsx
index 5c94350c87..f1eaed05ed 100644
--- a/packages/frontend/core/src/components/affine/ai-onboarding/index.tsx
+++ b/packages/frontend/core/src/components/affine/ai-onboarding/index.tsx
@@ -1,4 +1,9 @@
-import { FeatureFlagService, useService } from '@toeverything/infra';
+import { ServerConfigService } from '@affine/core/modules/cloud';
+import {
+  FeatureFlagService,
+  useLiveData,
+  useService,
+} from '@toeverything/infra';
 import { Suspense, useCallback, useEffect, useState } from 'react';
 
 import { AIOnboardingEdgeless } from './edgeless.dialog';
@@ -30,7 +35,12 @@ export const WorkspaceAIOnboarding = () => {
   const [dismissGeneral] = useDismiss(AIOnboardingType.GENERAL);
   const [dismissLocal] = useDismiss(AIOnboardingType.LOCAL);
   const featureFlagService = useService(FeatureFlagService);
-  const enableAI = featureFlagService.flags.enable_ai.value;
+  const serverConfigService = useService(ServerConfigService);
+  const serverFeatures = useLiveData(
+    serverConfigService.serverConfig.features$
+  );
+  const enableAI =
+    serverFeatures?.copilot && featureFlagService.flags.enable_ai.value;
 
   return (
     <Suspense>
@@ -43,7 +53,12 @@ export const WorkspaceAIOnboarding = () => {
 export const PageAIOnboarding = () => {
   const [dismissEdgeless] = useDismiss(AIOnboardingType.EDGELESS);
   const featureFlagService = useService(FeatureFlagService);
-  const enableAI = featureFlagService.flags.enable_ai.value;
+  const serverConfigService = useService(ServerConfigService);
+  const serverFeatures = useLiveData(
+    serverConfigService.serverConfig.features$
+  );
+  const enableAI =
+    serverFeatures?.copilot && featureFlagService.flags.enable_ai.value;
 
   return (
     <Suspense>
diff --git a/packages/frontend/core/src/components/affine/setting-modal/general-setting/editor/general.tsx b/packages/frontend/core/src/components/affine/setting-modal/general-setting/editor/general.tsx
index 7e341bf2a8..98e9494626 100644
--- a/packages/frontend/core/src/components/affine/setting-modal/general-setting/editor/general.tsx
+++ b/packages/frontend/core/src/components/affine/setting-modal/general-setting/editor/general.tsx
@@ -15,6 +15,7 @@ import {
   SettingRow,
   SettingWrapper,
 } from '@affine/component/setting-components';
+import { ServerConfigService } from '@affine/core/modules/cloud';
 import {
   EditorSettingService,
   type FontFamily,
@@ -402,8 +403,13 @@ export const SpellCheckSettings = () => {
 const AISettings = () => {
   const t = useI18n();
   const { openConfirmModal } = useConfirmModal();
-  const { featureFlagService } = useServices({ FeatureFlagService });
-
+  const { featureFlagService, serverConfigService } = useServices({
+    FeatureFlagService,
+    ServerConfigService,
+  });
+  const serverFeatures = useLiveData(
+    serverConfigService.serverConfig.features$
+  );
   const enableAI = useLiveData(featureFlagService.flags.enable_ai.$);
 
   const onAIChange = useCallback(
@@ -440,6 +446,10 @@ const AISettings = () => {
     [openConfirmModal, t, onAIChange]
   );
 
+  if (!serverFeatures?.copilot) {
+    return null;
+  }
+
   return (
     <SettingRow
       name={t['com.affine.settings.editorSettings.general.ai.title']()}
diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx b/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx
index 12d0bb6cf1..4017f2d99b 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/lit-adaper.tsx
@@ -4,6 +4,7 @@ import {
   useLitPortalFactory,
 } from '@affine/component';
 import { useJournalInfoHelper } from '@affine/core/components/hooks/use-journal';
+import { ServerConfigService } from '@affine/core/modules/cloud';
 import { EditorService } from '@affine/core/modules/editor';
 import { EditorSettingService } from '@affine/core/modules/editor-settting';
 import { toURLSearchParams } from '@affine/core/modules/navigation';
@@ -86,6 +87,7 @@ const usePatchSpecs = (shared: boolean, mode: DocMode) => {
     editorService,
     workspaceService,
     featureFlagService,
+    serverConfigService,
   } = useServices({
     PeekViewService,
     DocService,
@@ -93,8 +95,12 @@ const usePatchSpecs = (shared: boolean, mode: DocMode) => {
     WorkspaceService,
     EditorService,
     FeatureFlagService,
+    ServerConfigService,
   });
   const framework = useFramework();
+  const serverFeatures = useLiveData(
+    serverConfigService.serverConfig.features$
+  );
   const referenceRenderer: ReferenceReactRenderer = useMemo(() => {
     return function customReference(reference) {
       const data = reference.delta.attributes?.reference;
@@ -120,11 +126,17 @@ const usePatchSpecs = (shared: boolean, mode: DocMode) => {
   }, [workspaceService]);
 
   const specs = useMemo(() => {
-    const enableAI = featureFlagService.flags.enable_ai.value;
+    const enableAI =
+      serverFeatures?.copilot && featureFlagService.flags.enable_ai.value;
     return mode === 'edgeless'
-      ? createEdgelessModeSpecs(framework, enableAI)
-      : createPageModeSpecs(framework, enableAI);
-  }, [featureFlagService, mode, framework]);
+      ? createEdgelessModeSpecs(framework, !!enableAI)
+      : createPageModeSpecs(framework, !!enableAI);
+  }, [
+    serverFeatures?.copilot,
+    featureFlagService.flags.enable_ai.value,
+    mode,
+    framework,
+  ]);
 
   const confirmModal = useConfirmModal();
   const patchedSpecs = useMemo(() => {