feat(editor): feature flag store extension builder (#12235)

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->

## Summary by CodeRabbit

- **New Features**
  - Introduced feature flag synchronization for enhanced control over feature availability.
  - Added new configuration options for store management, allowing initialization and feature flag setup.

- **Improvements**
  - Updated how store extensions are accessed throughout the app for more robust initialization and configuration.
  - Enhanced workspace entities to support feature flag services, improving flexibility for workspace-specific features.
  - Centralized configuration of storage implementations for Electron environments.

- **Refactor**
  - Simplified editor module by removing direct feature flag synchronization logic.
  - Streamlined imports and configuration for storage modules, especially in Electron-based apps.

- **Bug Fixes**
  - Ensured consistent retrieval of store extensions across various modules and platforms.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/frontend/apps/android/src/app.tsx

@@ -252,7 +252,8 @@ framework.scope(ServerScope).override(AuthProvider, resolver => {
 
     const container = new Container();
     getStoreManager()
-      .get('store')
+      .config.init()
+      .value.get('store')
       .forEach(ext => {
         ext.setup(container);
       });

packages/frontend/apps/electron-renderer/src/app/effects/modules.ts

@@ -1,3 +1,4 @@
+import { configureElectronStateStorageImpls } from '@affine/core/desktop/storage';
 import { configureCommonModules } from '@affine/core/modules';
 import { configureAppTabsHeaderModule } from '@affine/core/modules/app-tabs-header';
 import { configureDesktopBackupModule } from '@affine/core/modules/backup';
@@ -11,7 +12,6 @@ import {
   configureTraySettingModule,
 } from '@affine/core/modules/editor-setting';
 import { configureFindInPageModule } from '@affine/core/modules/find-in-page';
-import { configureElectronStateStorageImpls } from '@affine/core/modules/storage';
 import {
   ClientSchemeProvider,
   PopupWindowProvider,

packages/frontend/apps/electron-renderer/src/popup/app.tsx

@@ -1,10 +1,8 @@
 import { ThemeProvider } from '@affine/core/components/theme-provider';
+import { configureElectronStateStorageImpls } from '@affine/core/desktop/storage';
 import { configureDesktopApiModule } from '@affine/core/modules/desktop-api';
 import { configureI18nModule, I18nProvider } from '@affine/core/modules/i18n';
-import {
-  configureElectronStateStorageImpls,
-  configureStorageModule,
-} from '@affine/core/modules/storage';
+import { configureStorageModule } from '@affine/core/modules/storage';
 import { configureEssentialThemeModule } from '@affine/core/modules/theme';
 import { appInfo } from '@affine/electron-api';
 import { Framework, FrameworkRoot } from '@toeverything/infra';

packages/frontend/apps/electron-renderer/src/shell/app.tsx

@@ -1,6 +1,7 @@
 import { useAppSettingHelper } from '@affine/core/components/hooks/affine/use-app-setting-helper';
 import { WindowsAppControls } from '@affine/core/components/pure/header/windows-app-controls';
 import { ThemeProvider } from '@affine/core/components/theme-provider';
+import { configureElectronStateStorageImpls } from '@affine/core/desktop/storage';
 import { configureAppSidebarModule } from '@affine/core/modules/app-sidebar';
 import { ShellAppSidebarFallback } from '@affine/core/modules/app-sidebar/views';
 import {
@@ -9,10 +10,7 @@ import {
 } from '@affine/core/modules/app-tabs-header';
 import { configureDesktopApiModule } from '@affine/core/modules/desktop-api';
 import { configureI18nModule, I18nProvider } from '@affine/core/modules/i18n';
-import {
-  configureElectronStateStorageImpls,
-  configureStorageModule,
-} from '@affine/core/modules/storage';
+import { configureStorageModule } from '@affine/core/modules/storage';
 import { configureAppThemeModule } from '@affine/core/modules/theme';
 import { Framework, FrameworkRoot } from '@toeverything/infra';
 

packages/frontend/apps/ios/src/app.tsx

@@ -264,7 +264,8 @@ const frameworkProvider = framework.provider();
 
     const container = new Container();
     getStoreManager()
-      .get('store')
+      .config.init()
+      .value.get('store')
       .forEach(ext => {
         ext.setup(container);
       });
@@ -309,7 +310,7 @@ const frameworkProvider = framework.provider();
       collection: workspace.docCollection,
       schema: getAFFiNEWorkspaceSchema(),
       markdown,
-      extensions: getStoreManager().get('store'),
+      extensions: getStoreManager().config.init().value.get('store'),
     });
     const docsService = workspace.scope.get(DocsService);
     if (docId) {

packages/frontend/core/src/blocksuite/ai/components/text-renderer.ts

@@ -252,7 +252,8 @@ export class TextRenderer extends WithDisposable(ShadowlessElement) {
       } else {
         const container = new Container();
         getStoreManager()
-          .get('store')
+          .config.init()
+          .value.get('store')
           .forEach(ext => {
             ext.setup(container);
           });

packages/frontend/core/src/blocksuite/ai/mini-mindmap/__tests__/mindmap-preview.unit.spec.ts

@@ -7,7 +7,7 @@ import { markdownToMindmap } from '../mindmap-preview.js';
 
 const container = new Container();
 getStoreManager()
-  .get('store')
+  .value.get('store')
   .forEach(ext => {
     ext.setup(container);
   });

packages/frontend/core/src/blocksuite/block-suite-page-list/utils.tsx

@@ -99,7 +99,7 @@ export const usePageHelper = (docCollection: Workspace) => {
       showImportModal({
         collection: docCollection,
         schema: getAFFiNEWorkspaceSchema(),
-        extensions: getStoreManager().get('store'),
+        extensions: getStoreManager().config.init().value.get('store'),
         onSuccess,
         onFail: message => {
           reject(new Error(message));

packages/frontend/core/src/blocksuite/extensions/feature-flag/feature-flag-syncer.ts

@@ -0,0 +1,29 @@
+import {
+  AFFINE_FLAGS,
+  type FeatureFlagService,
+} from '@affine/core/modules/feature-flag';
+import { FeatureFlagService as BSFeatureFlagService } from '@blocksuite/affine/shared/services';
+import { type ExtensionType, StoreExtension } from '@blocksuite/affine/store';
+
+export function getFeatureFlagSyncer(
+  featureFlagService: FeatureFlagService
+): ExtensionType {
+  class FeatureFlagSyncer extends StoreExtension {
+    static override key = 'feature-flag-syncer';
+
+    override loaded() {
+      const bsFeatureFlagService = this.store.get(BSFeatureFlagService);
+      Object.entries(AFFINE_FLAGS).forEach(([key, flag]) => {
+        if (flag.category === 'blocksuite') {
+          const value =
+            featureFlagService.flags[key as keyof AFFINE_FLAGS].value;
+          if (value !== undefined) {
+            bsFeatureFlagService.setFlag(flag.bsFlag, value);
+          }
+        }
+      });
+    }
+  }
+
+  return FeatureFlagSyncer;
+}

packages/frontend/core/src/blocksuite/extensions/feature-flag/index.ts

@@ -0,0 +1,29 @@
+import { getFeatureFlagSyncer } from '@affine/core/blocksuite/extensions/feature-flag/feature-flag-syncer';
+import { FeatureFlagService } from '@affine/core/modules/feature-flag';
+import {
+  type StoreExtensionContext,
+  StoreExtensionProvider,
+} from '@blocksuite/affine/ext-loader';
+import { z } from 'zod';
+
+const optionsSchema = z.object({
+  featureFlagService: z.instanceof(FeatureFlagService).optional(),
+});
+
+export class FeatureFlagStoreExtension extends StoreExtensionProvider {
+  override name = 'feature-flag-store-extension';
+
+  override schema = optionsSchema;
+
+  override setup(
+    context: StoreExtensionContext,
+    options?: z.infer<typeof optionsSchema>
+  ) {
+    super.setup(context, options);
+    const featureFlagService = options?.featureFlagService;
+    if (!featureFlagService) {
+      return;
+    }
+    context.register(getFeatureFlagSyncer(featureFlagService));
+  }
+}

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

@@ -1,3 +1,5 @@
+import { FeatureFlagStoreExtension } from '@affine/core/blocksuite/extensions/feature-flag';
+import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import {
   type StoreExtensionContext,
   StoreExtensionManager,
@@ -18,11 +20,57 @@ class MigratingAffineStoreExtension extends StoreExtensionProvider {
   }
 }
 
-const manager = new StoreExtensionManager([
-  ...getInternalStoreExtensions(),
-  MigratingAffineStoreExtension,
-]);
+interface Configure {
+  init: () => Configure;
+  featureFlag: (featureFlagService?: FeatureFlagService) => Configure;
+  value: StoreExtensionManager;
+}
+
+class StoreProvider {
+  static instance: StoreProvider | null = null;
+  static getInstance() {
+    if (!StoreProvider.instance) {
+      StoreProvider.instance = new StoreProvider();
+    }
+    return StoreProvider.instance;
+  }
+
+  private readonly _manager: StoreExtensionManager;
+
+  constructor() {
+    this._manager = new StoreExtensionManager([
+      ...getInternalStoreExtensions(),
+      MigratingAffineStoreExtension,
+      FeatureFlagStoreExtension,
+    ]);
+  }
+
+  get config(): Configure {
+    return {
+      init: this._initDefaultConfig,
+      featureFlag: this._configureFeatureFlag,
+      value: this._manager,
+    };
+  }
+
+  get value(): StoreExtensionManager {
+    return this._manager;
+  }
+
+  private readonly _initDefaultConfig = () => {
+    this.config.featureFlag();
+
+    return this.config;
+  };
+
+  private readonly _configureFeatureFlag = (
+    featureFlagService?: FeatureFlagService
+  ) => {
+    this._manager.configure(FeatureFlagStoreExtension, { featureFlagService });
+    return this.config;
+  };
+}
 
 export function getStoreManager() {
-  return manager;
+  return StoreProvider.getInstance();
 }

packages/frontend/core/src/components/page-list/__tests__/use-block-suite-page-preview.spec.ts

@@ -13,7 +13,7 @@ import { beforeEach, describe, expect, test, vi } from 'vitest';
 import { useBlockSuitePagePreview } from '../use-block-suite-page-preview';
 let docCollection: TestWorkspace;
 
-const extensions = getStoreManager().get('store');
+const extensions = getStoreManager().config.init().value.get('store');
 
 beforeEach(async () => {
   vi.useFakeTimers({ toFake: ['requestIdleCallback'] });

packages/frontend/core/src/desktop/dialogs/import/index.tsx

@@ -146,7 +146,7 @@ const importConfigs: Record<ImportType, ImportConfig> = {
           schema: getAFFiNEWorkspaceSchema(),
           markdown: text,
           fileName,
-          extensions: getStoreManager().get('store'),
+          extensions: getStoreManager().config.init().value.get('store'),
         });
         if (docId) docIds.push(docId);
       }
@@ -165,7 +165,7 @@ const importConfigs: Record<ImportType, ImportConfig> = {
         collection: docCollection,
         schema: getAFFiNEWorkspaceSchema(),
         imported: file,
-        extensions: getStoreManager().get('store'),
+        extensions: getStoreManager().config.init().value.get('store'),
       });
       return {
         docIds,
@@ -185,7 +185,7 @@ const importConfigs: Record<ImportType, ImportConfig> = {
         const docId = await HtmlTransformer.importHTMLToDoc({
           collection: docCollection,
           schema: getAFFiNEWorkspaceSchema(),
-          extensions: getStoreManager().get('store'),
+          extensions: getStoreManager().config.init().value.get('store'),
           html: text,
           fileName,
         });
@@ -207,7 +207,7 @@ const importConfigs: Record<ImportType, ImportConfig> = {
           collection: docCollection,
           schema: getAFFiNEWorkspaceSchema(),
           imported: file,
-          extensions: getStoreManager().get('store'),
+          extensions: getStoreManager().config.init().value.get('store'),
         });
       return {
         docIds: pageIds,

packages/frontend/core/src/desktop/storage.ts

@@ -0,0 +1,18 @@
+import { DesktopApiService } from '@affine/core/modules/desktop-api';
+import {
+  CacheStorage,
+  GlobalCache,
+  GlobalState,
+} from '@affine/core/modules/storage';
+import {
+  ElectronGlobalCache,
+  ElectronGlobalState,
+} from '@affine/core/modules/storage/impls/electron';
+import { IDBGlobalState } from '@affine/core/modules/storage/impls/storage';
+import type { Framework } from '@toeverything/infra';
+
+export function configureElectronStateStorageImpls(framework: Framework) {
+  framework.impl(GlobalCache, ElectronGlobalCache, [DesktopApiService]);
+  framework.impl(GlobalState, ElectronGlobalState, [DesktopApiService]);
+  framework.impl(CacheStorage, IDBGlobalState);
+}

packages/frontend/core/src/modules/editor/entities/editor.ts

@@ -5,10 +5,7 @@ import { DefaultTool } from '@blocksuite/affine/blocks/surface';
 import type { DocTitle } from '@blocksuite/affine/fragments/doc-title';
 import type { DocMode, ReferenceParams } from '@blocksuite/affine/model';
 import { HighlightSelection } from '@blocksuite/affine/shared/selection';
-import {
-  DocModeProvider,
-  FeatureFlagService as BSFeatureFlagService,
-} from '@blocksuite/affine/shared/services';
+import { DocModeProvider } from '@blocksuite/affine/shared/services';
 import { GfxControllerIdentifier } from '@blocksuite/affine/std/gfx';
 import type { InlineEditor } from '@blocksuite/std/inline';
 import { effect } from '@preact/signals-core';
@@ -17,7 +14,6 @@ import { defaults, isEqual, omit } from 'lodash-es';
 import { skip } from 'rxjs';
 
 import type { DocService } from '../../doc';
-import { AFFINE_FLAGS, type FeatureFlagService } from '../../feature-flag';
 import { paramsParseOptions, preprocessParams } from '../../navigation/utils';
 import type { WorkbenchView } from '../../workbench';
 import type { WorkspaceService } from '../../workspace';
@@ -196,7 +192,6 @@ export class Editor extends Entity {
       throw new Error('already bound');
     }
 
-    this._setupBlocksuiteEditorFlags(editorContainer);
     this.editorContainer$.next(editorContainer);
     const unsubs: (() => void)[] = [];
 
@@ -325,24 +320,9 @@ export class Editor extends Entity {
     };
   }
 
-  private _setupBlocksuiteEditorFlags(editorContainer: AffineEditorContainer) {
-    const affineFeatureFlagService = this.featureFlagService;
-    const bsFeatureFlagService = editorContainer.doc.get(BSFeatureFlagService);
-    Object.entries(AFFINE_FLAGS).forEach(([key, flag]) => {
-      if (flag.category === 'blocksuite') {
-        const value =
-          affineFeatureFlagService.flags[key as keyof AFFINE_FLAGS].value;
-        if (value !== undefined) {
-          bsFeatureFlagService.setFlag(flag.bsFlag, value);
-        }
-      }
-    });
-  }
-
   constructor(
     private readonly docService: DocService,
-    private readonly workspaceService: WorkspaceService,
-    private readonly featureFlagService: FeatureFlagService
+    private readonly workspaceService: WorkspaceService
   ) {
     super();
   }

packages/frontend/core/src/modules/editor/index.ts

@@ -1,7 +1,6 @@
 import { type Framework } from '@toeverything/infra';
 
 import { DocScope, DocService } from '../doc';
-import { FeatureFlagService } from '../feature-flag';
 import { WorkspaceScope, WorkspaceService } from '../workspace';
 import { Editor } from './entities/editor';
 import { EditorScope } from './scopes/editor';
@@ -19,7 +18,7 @@ export function configureEditorModule(framework: Framework) {
     .scope(WorkspaceScope)
     .scope(DocScope)
     .service(EditorsService)
-    .entity(Editor, [DocService, WorkspaceService, FeatureFlagService])
+    .entity(Editor, [DocService, WorkspaceService])
     .scope(EditorScope)
     .service(EditorService, [EditorScope]);
 }

packages/frontend/core/src/modules/import-clipper/services/import.ts

@@ -35,7 +35,7 @@ export class ImportClipperService extends Service {
       collection: workspace.docCollection,
       schema: getAFFiNEWorkspaceSchema(),
       markdown: clipperInput.contentMarkdown,
-      extensions: getStoreManager().get('store'),
+      extensions: getStoreManager().config.init().value.get('store'),
     });
     const docsService = workspace.scope.get(DocsService);
     if (docId) {
@@ -69,7 +69,7 @@ export class ImportClipperService extends Service {
           collection: docCollection,
           schema: getAFFiNEWorkspaceSchema(),
           markdown: clipperInput.contentMarkdown,
-          extensions: getStoreManager().get('store'),
+          extensions: getStoreManager().config.init().value.get('store'),
         });
       }
     );

packages/frontend/core/src/modules/integration/entities/writer.ts

@@ -61,7 +61,7 @@ export class IntegrationWriter extends Entity {
         schema: getAFFiNEWorkspaceSchema(),
         markdown,
         fileName: title,
-        extensions: getStoreManager().get('store'),
+        extensions: getStoreManager().config.init().value.get('store'),
       });
 
       if (!newDocId) throw new Error('Failed to create a new doc');
@@ -85,7 +85,7 @@ export class IntegrationWriter extends Entity {
           doc,
           blockId: noteBlockId,
           markdown,
-          extensions: getStoreManager().get('store'),
+          extensions: getStoreManager().config.init().value.get('store'),
         });
       } else if (updateStrategy === 'append') {
         const pageBlockId = doc.getBlocksByFlavour('affine:page')[0]?.id;
@@ -94,7 +94,7 @@ export class IntegrationWriter extends Entity {
           doc,
           blockId,
           markdown: `---\n${markdown}`,
-          extensions: getStoreManager().get('store'),
+          extensions: getStoreManager().config.init().value.get('store'),
         });
       } else {
         throw new Error('Invalid update strategy');

packages/frontend/core/src/modules/storage/index.ts

@@ -14,8 +14,6 @@ export { NbstoreService } from './services/nbstore';
 
 import { type Framework } from '@toeverything/infra';
 
-import { DesktopApiService } from '../desktop-api';
-import { ElectronGlobalCache, ElectronGlobalState } from './impls/electron';
 import {
   IDBGlobalState,
   LocalStorageGlobalCache,
@@ -49,12 +47,6 @@ export function configureLocalStorageStateStorageImpls(framework: Framework) {
   framework.impl(CacheStorage, IDBGlobalState);
 }
 
-export function configureElectronStateStorageImpls(framework: Framework) {
-  framework.impl(GlobalCache, ElectronGlobalCache, [DesktopApiService]);
-  framework.impl(GlobalState, ElectronGlobalState, [DesktopApiService]);
-  framework.impl(CacheStorage, IDBGlobalState);
-}
-
 export function configureCommonGlobalStorageImpls(framework: Framework) {
   framework.impl(GlobalSessionState, SessionStorageGlobalSessionState);
 }

packages/frontend/core/src/modules/workspace/entities/workspace.ts

@@ -1,3 +1,4 @@
+import type { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import type { Workspace as WorkspaceInterface } from '@blocksuite/affine/store';
 import { Entity, LiveData, yjsGetPath } from '@toeverything/infra';
 import type { Observable } from 'rxjs';
@@ -9,7 +10,10 @@ import type { WorkspaceScope } from '../scopes/workspace';
 import { WorkspaceEngineService } from '../services/engine';
 
 export class Workspace extends Entity {
-  constructor(public readonly scope: WorkspaceScope) {
+  constructor(
+    public readonly scope: WorkspaceScope,
+    public readonly featureFlagService: FeatureFlagService
+  ) {
     super();
   }
 
@@ -30,6 +34,7 @@ export class Workspace extends Entity {
       this._docCollection = new WorkspaceImpl({
         id: this.openOptions.metadata.id,
         rootDoc: this.rootYDoc,
+        featureFlagService: this.featureFlagService,
         blobSource: {
           get: async key => {
             const record = await this.engine.blob.get(key);

packages/frontend/core/src/modules/workspace/impls/doc.ts

@@ -5,20 +5,21 @@ import {
   type ExtensionType,
   type GetStoreOptions,
   StoreContainer,
-  type Workspace,
   type YBlock,
 } from '@blocksuite/affine/store';
 import { Awareness } from 'y-protocols/awareness.js';
 import * as Y from 'yjs';
 
+import type { WorkspaceImpl } from './workspace';
+
 type DocOptions = {
   id: string;
-  collection: Workspace;
+  collection: WorkspaceImpl;
   doc: Y.Doc;
 };
 
 export class DocImpl implements Doc {
-  private readonly _collection: Workspace;
+  private readonly _collection: WorkspaceImpl;
 
   private readonly _storeContainer: StoreContainer;
 
@@ -136,7 +137,10 @@ export class DocImpl implements Doc {
     extensions,
     id,
   }: GetStoreOptions = {}) {
-    const storeExtensions = getStoreManager().get('store');
+    const storeExtensions = getStoreManager()
+      .config.init()
+      .featureFlag(this.workspace.featureFlagService)
+      .value.get('store');
     const exts = storeExtensions
       .concat(extensions ?? [])
       .concat(this.storeExtensions);

packages/frontend/core/src/modules/workspace/impls/workspace.ts

@@ -19,6 +19,7 @@ import { Subject } from 'rxjs';
 import type { Awareness } from 'y-protocols/awareness.js';
 import type { Doc as YDoc } from 'yjs';
 
+import type { FeatureFlagService } from '../../feature-flag';
 import { DocImpl } from './doc';
 import { WorkspaceMetaImpl } from './meta';
 
@@ -29,6 +30,7 @@ type WorkspaceOptions = {
   onLoadDoc?: (doc: YDoc) => void;
   onLoadAwareness?: (awareness: Awareness) => void;
   onCreateDoc?: (docId?: string) => string;
+  featureFlagService?: FeatureFlagService;
 };
 
 export class WorkspaceImpl implements Workspace {
@@ -57,6 +59,7 @@ export class WorkspaceImpl implements Workspace {
   readonly onLoadDoc?: (doc: YDoc) => void;
   readonly onLoadAwareness?: (awareness: Awareness) => void;
   readonly onCreateDoc?: (docId?: string) => string;
+  readonly featureFlagService?: FeatureFlagService;
 
   constructor({
     id,
@@ -65,8 +68,10 @@ export class WorkspaceImpl implements Workspace {
     onLoadDoc,
     onLoadAwareness,
     onCreateDoc,
+    featureFlagService,
   }: WorkspaceOptions) {
     this.id = id || '';
+    this.featureFlagService = featureFlagService;
     this.doc = rootDoc;
     this.onLoadDoc = onLoadDoc;
     this.onLoadDoc?.(this.doc);

packages/frontend/core/src/modules/workspace/index.ts

@@ -1,3 +1,5 @@
+import { FeatureFlagService } from '@affine/core/modules/feature-flag';
+
 export type { WorkspaceProfileInfo } from './entities/profile';
 export { Workspace } from './entities/workspace';
 export { WorkspaceEngineBeforeStart, WorkspaceInitialized } from './events';
@@ -70,7 +72,7 @@ export function configureWorkspaceModule(framework: Framework) {
     ])
     .scope(WorkspaceScope)
     .service(WorkspaceService)
-    .entity(Workspace, [WorkspaceScope])
+    .entity(Workspace, [WorkspaceScope, FeatureFlagService])
     .service(WorkspaceEngineService, [WorkspaceScope])
     .entity(WorkspaceEngine, [WorkspaceService, NbstoreService])
     .impl(WorkspaceLocalState, WorkspaceLocalStateImpl, [