feat: run app in closure (#3790)

apps/core/src/_plugin/index.test.tsx

@@ -1,17 +1,23 @@
 import { assertExists } from '@blocksuite/global/utils';
-import { loadedPluginNameAtom, rootStore } from '@toeverything/infra/atom';
+import {
+  getCurrentStore,
+  loadedPluginNameAtom,
+} from '@toeverything/infra/atom';
 import { use } from 'foxact/use';
 import { useAtomValue } from 'jotai';
 import { Provider } from 'jotai/react';
 import { StrictMode } from 'react';
 import { createRoot } from 'react-dom/client';
 
-import { _pluginNestedImportsMap } from '../bootstrap/plugins/setup';
+import { createSetup } from '../bootstrap/plugins/setup';
-import { pluginRegisterPromise } from '../bootstrap/register-plugins';
+import { bootstrapPluginSystem } from '../bootstrap/register-plugins';
 
 async function main() {
   const { setup } = await import('../bootstrap/setup');
-  await setup();
+  const rootStore = getCurrentStore();
+  await setup(rootStore);
+  const { _pluginNestedImportsMap } = createSetup(rootStore);
+  const pluginRegisterPromise = bootstrapPluginSystem(rootStore);
   const root = document.getElementById('app');
   assertExists(root);
 

apps/core/src/adapters/local/index.tsx

@@ -20,7 +20,7 @@ import { getOrCreateWorkspace } from '@affine/workspace/manager';
 import { createIndexedDBDownloadProvider } from '@affine/workspace/providers';
 import { nanoid } from '@blocksuite/store';
 import { useStaticBlockSuiteWorkspace } from '@toeverything/infra/__internal__/react';
-import { rootStore } from '@toeverything/infra/atom';
+import { getCurrentStore } from '@toeverything/infra/atom';
 import { buildShowcaseWorkspace } from '@toeverything/infra/blocksuite';
 
 import { setPageModeAtom } from '../../atoms';
@@ -46,7 +46,7 @@ export const LocalAdapter: WorkspaceAdapter<WorkspaceFlavour.LOCAL> = {
       blockSuiteWorkspace.meta.setName(DEFAULT_WORKSPACE_NAME);
       if (runtimeConfig.enablePreloading) {
         buildShowcaseWorkspace(blockSuiteWorkspace, {
-          store: rootStore,
+          store: getCurrentStore(),
           atoms: {
             pageMode: setPageModeAtom,
           },

apps/core/src/app.tsx

@@ -5,6 +5,7 @@ import '@toeverything/components/style.css';
 import { AffineContext } from '@affine/component/context';
 import { WorkspaceFallback } from '@affine/component/workspace';
 import { CacheProvider } from '@emotion/react';
+import { getCurrentStore } from '@toeverything/infra/atom';
 import { use } from 'foxact/use';
 import type { PropsWithChildren, ReactElement } from 'react';
 import { lazy, memo, Suspense } from 'react';
@@ -47,7 +48,7 @@ export const App = memo(function App() {
   use(languageLoadingPromise);
   return (
     <CacheProvider value={cache}>
-      <AffineContext>
+      <AffineContext store={getCurrentStore()}>
         <DebugProvider>
           <RouterProvider
             fallbackElement={<WorkspaceFallback key="RouterFallback" />}

apps/core/src/bootstrap/plugins/setup.ts

@@ -18,10 +18,10 @@ import {
   contentLayoutAtom,
   currentPageAtom,
   currentWorkspaceAtom,
-  rootStore,
 } from '@toeverything/infra/atom';
 import { atom } from 'jotai';
 import { Provider } from 'jotai/react';
+import type { createStore } from 'jotai/vanilla';
 import { createElement, type PropsWithChildren } from 'react';
 
 import { createFetch } from './endowments/fercher';
@@ -32,6 +32,7 @@ const dynamicImportKey = '$h_import';
 
 const permissionLogger = new DebugLogger('plugins:permission');
 const importLogger = new DebugLogger('plugins:import');
+const entryLogger = new DebugLogger('plugins:entry');
 
 const pushLayoutAtom = atom<
   null,
@@ -39,7 +40,11 @@ const pushLayoutAtom = atom<
   [
     pluginName: string,
     create: (root: HTMLElement) => () => void,
-    options: { maxWidth: (number | undefined)[] } | undefined,
+    options:
+      | {
+          maxWidth: (number | undefined)[];
+        }
+      | undefined,
   ],
   void
 >(null, (_, set, pluginName, callback, options) => {
@@ -106,438 +111,469 @@ const deleteLayoutAtom = atom<null, [string], void>(null, (_, set, id) => {
   });
 });
 
-// module -> importName -> updater[]
+const setupWeakMap = new WeakMap<
-export const _rootImportsMap = new Map<string, Map<string, any>>();
+  ReturnType<typeof createStore>,
-const rootImportsMapSetupPromise = setupImportsMap(_rootImportsMap, {
+  ReturnType<typeof createSetupImpl>
-  react: import('react'),
-  'react/jsx-runtime': import('react/jsx-runtime'),
-  'react-dom': import('react-dom'),
-  'react-dom/client': import('react-dom/client'),
-  jotai: import('jotai'),
-  'jotai/utils': import('jotai/utils'),
-  swr: import('swr'),
-  '@affine/component': import('@affine/component'),
-  '@blocksuite/icons': import('@blocksuite/icons'),
-  '@blocksuite/blocks': import('@blocksuite/blocks'),
-  '@affine/sdk/entry': {
-    rootStore: rootStore,
-    currentWorkspaceAtom: currentWorkspaceAtom,
-    currentPageAtom: currentPageAtom,
-    pushLayoutAtom: pushLayoutAtom,
-    deleteLayoutAtom: deleteLayoutAtom,
-  },
-  '@blocksuite/global/utils': import('@blocksuite/global/utils'),
-  '@toeverything/infra/atom': import('@toeverything/infra/atom'),
-  '@toeverything/components/button': import('@toeverything/components/button'),
-});
-
-// pluginName -> module -> importName -> updater[]
-export const _pluginNestedImportsMap = new Map<
-  string,
-  Map<string, Map<string, any>>
 >();
 
-const pluginImportsFunctionMap = new Map<string, (imports: any) => void>();
+export function createSetup(rootStore: ReturnType<typeof createStore>) {
-export const createImports = (pluginName: string) => {
+  if (setupWeakMap.has(rootStore)) {
-  if (pluginImportsFunctionMap.has(pluginName)) {
     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-    return pluginImportsFunctionMap.get(pluginName)!;
+    return setupWeakMap.get(rootStore)!;
   }
-  const imports = (
+  const setup = createSetupImpl(rootStore);
-    newUpdaters: [string, [string, ((val: any) => void)[]][]][]
+  setupWeakMap.set(rootStore, setup);
+  return setup;
+}
+
+function createSetupImpl(rootStore: ReturnType<typeof createStore>) {
+  // clean up plugin windows when switching to other pages
+  rootStore.sub(currentPageAtom, () => {
+    rootStore.set(contentLayoutAtom, 'editor');
+  });
+
+  // module -> importName -> updater[]
+  const _rootImportsMap = new Map<string, Map<string, any>>();
+  const rootImportsMapSetupPromise = setupImportsMap(_rootImportsMap, {
+    react: import('react'),
+    'react/jsx-runtime': import('react/jsx-runtime'),
+    'react-dom': import('react-dom'),
+    'react-dom/client': import('react-dom/client'),
+    jotai: import('jotai'),
+    'jotai/utils': import('jotai/utils'),
+    swr: import('swr'),
+    '@affine/component': import('@affine/component'),
+    '@blocksuite/icons': import('@blocksuite/icons'),
+    '@blocksuite/blocks': import('@blocksuite/blocks'),
+    '@affine/sdk/entry': {
+      rootStore,
+      currentWorkspaceAtom: currentWorkspaceAtom,
+      currentPageAtom: currentPageAtom,
+      pushLayoutAtom: pushLayoutAtom,
+      deleteLayoutAtom: deleteLayoutAtom,
+    },
+    '@blocksuite/global/utils': import('@blocksuite/global/utils'),
+    '@toeverything/infra/atom': import('@toeverything/infra/atom'),
+    '@toeverything/components/button': import(
+      '@toeverything/components/button'
+    ),
+  });
+
+  // pluginName -> module -> importName -> updater[]
+  const _pluginNestedImportsMap = new Map<
+    string,
+    Map<string, Map<string, any>>
+  >();
+
+  const pluginImportsFunctionMap = new Map<string, (imports: any) => void>();
+  const createImports = (pluginName: string) => {
+    if (pluginImportsFunctionMap.has(pluginName)) {
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      return pluginImportsFunctionMap.get(pluginName)!;
+    }
+    const imports = (
+      newUpdaters: [string, [string, ((val: any) => void)[]][]][]
+    ) => {
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      const currentImportMap = _pluginNestedImportsMap.get(pluginName)!;
+      importLogger.debug('currentImportMap', pluginName, currentImportMap);
+
+      for (const [module, moduleUpdaters] of newUpdaters) {
+        importLogger.debug('imports module', module, moduleUpdaters);
+        let moduleImports = _rootImportsMap.get(module);
+        if (!moduleImports) {
+          moduleImports = currentImportMap.get(module);
+        }
+        if (moduleImports) {
+          for (const [importName, importUpdaters] of moduleUpdaters) {
+            const updateImport = (value: any) => {
+              for (const importUpdater of importUpdaters) {
+                importUpdater(value);
+              }
+            };
+            if (moduleImports.has(importName)) {
+              const val = moduleImports.get(importName);
+              updateImport(val);
+            }
+          }
+        } else {
+          console.error(
+            'cannot find module in plugin import map',
+            module,
+            currentImportMap,
+            _pluginNestedImportsMap
+          );
+        }
+      }
+    };
+    pluginImportsFunctionMap.set(pluginName, imports);
+    return imports;
+  };
+
+  const abortController = new AbortController();
+
+  const pluginFetch = createFetch({});
+  const timer = createTimers(abortController.signal);
+
+  const sharedGlobalThis = Object.assign(Object.create(null), timer, {
+    Object: globalThis.Object,
+    fetch: pluginFetch,
+    ReadableStream: globalThis.ReadableStream,
+    Symbol: globalThis.Symbol,
+    Error: globalThis.Error,
+    TypeError: globalThis.TypeError,
+    RangeError: globalThis.RangeError,
+    console: globalThis.console,
+    crypto: globalThis.crypto,
+  });
+
+  const dynamicImportMap = new Map<
+    string,
+    (moduleName: string) => Promise<any>
+  >();
+
+  const createOrGetDynamicImport = (baseUrl: string, pluginName: string) => {
+    if (dynamicImportMap.has(pluginName)) {
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      return dynamicImportMap.get(pluginName)!;
+    }
+    const dynamicImport = async (moduleName: string): Promise<any> => {
+      const codeUrl = `${baseUrl}/${moduleName}`;
+      const analysisUrl = `${baseUrl}/${moduleName}.json`;
+      const response = await fetch(codeUrl);
+      const analysisResponse = await fetch(analysisUrl);
+      const analysis = await analysisResponse.json();
+      const exports = analysis.exports as string[];
+      const code = await response.text();
+      const moduleCompartment = new Compartment(
+        createOrGetGlobalThis(
+          pluginName,
+          // use singleton here to avoid infinite loop
+          createOrGetDynamicImport(pluginName, baseUrl)
+        )
+      );
+      const entryPoint = moduleCompartment.evaluate(code, {
+        __evadeHtmlCommentTest__: true,
+      });
+      const moduleExports = {} as Record<string, any>;
+      const setVarProxy = new Proxy(
+        {},
+        {
+          get(_, p: string): any {
+            return (newValue: any) => {
+              moduleExports[p] = newValue;
+            };
+          },
+        }
+      );
+      entryPoint({
+        imports: createImports(pluginName),
+        liveVar: setVarProxy,
+        onceVar: setVarProxy,
+      });
+      importLogger.debug('import', moduleName, exports, moduleExports);
+      return moduleExports;
+    };
+    dynamicImportMap.set(pluginName, dynamicImport);
+    return dynamicImport;
+  };
+
+  const globalThisMap = new Map<string, any>();
+
+  const createOrGetGlobalThis = (
+    pluginName: string,
+    dynamicImport: (moduleName: string) => Promise<any>
   ) => {
+    if (globalThisMap.has(pluginName)) {
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      return globalThisMap.get(pluginName)!;
+    }
+    const pluginGlobalThis = Object.assign(
+      Object.create(null),
+      sharedGlobalThis,
+      {
+        // fixme: vite build output bundle will have this, we should remove it
+        process: Object.freeze({
+          env: {
+            NODE_ENV: process.env.NODE_ENV,
+          },
+        }),
+        // dynamic import function
+        [dynamicImportKey]: dynamicImport,
+        // UNSAFE: React will read `window` and `document`
+        window: new Proxy(
+          {},
+          {
+            get(_, key) {
+              permissionLogger.debug(`${pluginName} is accessing window`, key);
+              if (sharedGlobalThis[key]) return sharedGlobalThis[key];
+              const result = Reflect.get(window, key);
+              if (typeof result === 'function') {
+                if (result === ShadowRoot) {
+                  return result;
+                }
+                return function (...args: any[]) {
+                  permissionLogger.debug(
+                    `${pluginName} is calling window`,
+                    key,
+                    args
+                  );
+                  return result.apply(window, args);
+                };
+              }
+              permissionLogger.debug('window', key, result);
+              return result;
+            },
+          }
+        ),
+        document: new Proxy(
+          {},
+          {
+            get(_, key) {
+              permissionLogger.debug(
+                `${pluginName} is accessing document`,
+                key
+              );
+              if (sharedGlobalThis[key]) return sharedGlobalThis[key];
+              const result = Reflect.get(document, key);
+              if (typeof result === 'function') {
+                return function (...args: any[]) {
+                  permissionLogger.debug(
+                    `${pluginName} is calling window`,
+                    key,
+                    args
+                  );
+                  return result.apply(document, args);
+                };
+              }
+              permissionLogger.debug('document', key, result);
+              return result;
+            },
+          }
+        ),
+        navigator: {
+          userAgent: navigator.userAgent,
+        },
+
+        MouseEvent: globalThis.MouseEvent,
+        KeyboardEvent: globalThis.KeyboardEvent,
+        CustomEvent: globalThis.CustomEvent,
+
+        // copilot uses these
+        Date: globalThis.Date,
+        Math: globalThis.Math,
+        URL: globalThis.URL,
+        URLSearchParams: globalThis.URLSearchParams,
+        Headers: globalThis.Headers,
+        TextEncoder: globalThis.TextEncoder,
+        TextDecoder: globalThis.TextDecoder,
+        Request: globalThis.Request,
+
+        // image-preview uses these
+        Blob: globalThis.Blob,
+        ClipboardItem: globalThis.ClipboardItem,
+
+        // vue uses these
+        Element: globalThis.Element,
+        SVGElement: globalThis.SVGElement,
+
+        // fixme: use our own db api
+        indexedDB: globalThis.indexedDB,
+        IDBRequest: globalThis.IDBRequest,
+        IDBDatabase: globalThis.IDBDatabase,
+        IDBCursorWithValue: globalThis.IDBCursorWithValue,
+        IDBFactory: globalThis.IDBFactory,
+        IDBKeyRange: globalThis.IDBKeyRange,
+        IDBOpenDBRequest: globalThis.IDBOpenDBRequest,
+        IDBTransaction: globalThis.IDBTransaction,
+        IDBObjectStore: globalThis.IDBObjectStore,
+        IDBIndex: globalThis.IDBIndex,
+        IDBCursor: globalThis.IDBCursor,
+        IDBVersionChangeEvent: globalThis.IDBVersionChangeEvent,
+      }
+    );
+    pluginGlobalThis.global = pluginGlobalThis;
+    globalThisMap.set(pluginName, pluginGlobalThis);
+    return pluginGlobalThis;
+  };
+
+  const setupPluginCode = async (
+    baseUrl: string,
+    pluginName: string,
+    filename: string
+  ) => {
+    await rootImportsMapSetupPromise;
+    if (!_pluginNestedImportsMap.has(pluginName)) {
+      _pluginNestedImportsMap.set(pluginName, new Map());
+    }
     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
     const currentImportMap = _pluginNestedImportsMap.get(pluginName)!;
-    importLogger.debug('currentImportMap', pluginName, currentImportMap);
+    const isMissingPackage = (name: string) =>
+      _rootImportsMap.has(name) && !currentImportMap.has(name);
 
-    for (const [module, moduleUpdaters] of newUpdaters) {
+    const bundleAnalysis = await fetch(`${baseUrl}/${filename}.json`).then(
-      importLogger.debug('imports module', module, moduleUpdaters);
+      res => res.json()
-      let moduleImports = _rootImportsMap.get(module);
+    );
-      if (!moduleImports) {
+    const moduleExports = bundleAnalysis.exports as Record<string, [string]>;
-        moduleImports = currentImportMap.get(module);
+    const moduleImports = bundleAnalysis.imports as string[];
-      }
+    const moduleReexports = bundleAnalysis.reexports as Record<
-      if (moduleImports) {
+      string,
-        for (const [importName, importUpdaters] of moduleUpdaters) {
+      [localName: string, exportedName: string][]
-          const updateImport = (value: any) => {
+    >;
-            for (const importUpdater of importUpdaters) {
+    await Promise.all(
-              importUpdater(value);
+      moduleImports.map(name => {
-            }
+        if (isMissingPackage(name)) {
-          };
+          return Promise.resolve();
-          if (moduleImports.has(importName)) {
+        } else {
-            const val = moduleImports.get(importName);
+          importLogger.debug('missing package', name);
-            updateImport(val);
+          return setupPluginCode(baseUrl, pluginName, name);
-          }
         }
-      } else {
+      })
-        console.error(
+    );
-          'cannot find module in plugin import map',
+    const code = await fetch(
-          module,
+      `${baseUrl}/${filename.replace(/^\.\//, '')}`
-          currentImportMap,
+    ).then(res => res.text());
-          _pluginNestedImportsMap
+    importLogger.debug('evaluating', filename);
-        );
-      }
-    }
-  };
-  pluginImportsFunctionMap.set(pluginName, imports);
-  return imports;
-};
-
-const abortController = new AbortController();
-
-const pluginFetch = createFetch({});
-const timer = createTimers(abortController.signal);
-
-const sharedGlobalThis = Object.assign(Object.create(null), timer, {
-  Object: globalThis.Object,
-  fetch: pluginFetch,
-  ReadableStream: globalThis.ReadableStream,
-  Symbol: globalThis.Symbol,
-  Error: globalThis.Error,
-  TypeError: globalThis.TypeError,
-  RangeError: globalThis.RangeError,
-  console: globalThis.console,
-  crypto: globalThis.crypto,
-});
-
-const dynamicImportMap = new Map<
-  string,
-  (moduleName: string) => Promise<any>
->();
-
-export const createOrGetDynamicImport = (
-  baseUrl: string,
-  pluginName: string
-) => {
-  if (dynamicImportMap.has(pluginName)) {
-    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-    return dynamicImportMap.get(pluginName)!;
-  }
-  const dynamicImport = async (moduleName: string): Promise<any> => {
-    const codeUrl = `${baseUrl}/${moduleName}`;
-    const analysisUrl = `${baseUrl}/${moduleName}.json`;
-    const response = await fetch(codeUrl);
-    const analysisResponse = await fetch(analysisUrl);
-    const analysis = await analysisResponse.json();
-    const exports = analysis.exports as string[];
-    const code = await response.text();
     const moduleCompartment = new Compartment(
       createOrGetGlobalThis(
         pluginName,
         // use singleton here to avoid infinite loop
-        createOrGetDynamicImport(pluginName, baseUrl)
+        createOrGetDynamicImport(baseUrl, pluginName)
       )
     );
     const entryPoint = moduleCompartment.evaluate(code, {
       __evadeHtmlCommentTest__: true,
     });
-    const moduleExports = {} as Record<string, any>;
+    const moduleExportsMap = new Map<string, any>();
     const setVarProxy = new Proxy(
       {},
       {
         get(_, p: string): any {
           return (newValue: any) => {
-            moduleExports[p] = newValue;
+            moduleExportsMap.set(p, newValue);
           };
         },
       }
     );
+    currentImportMap.set(filename, moduleExportsMap);
     entryPoint({
       imports: createImports(pluginName),
       liveVar: setVarProxy,
       onceVar: setVarProxy,
     });
-    importLogger.debug('import', moduleName, exports, moduleExports);
+
-    return moduleExports;
+    for (const [newExport, [originalExport]] of Object.entries(moduleExports)) {
+      if (newExport === originalExport) continue;
+      const value = moduleExportsMap.get(originalExport);
+      moduleExportsMap.set(newExport, value);
+      moduleExportsMap.delete(originalExport);
+    }
+
+    for (const [name, reexports] of Object.entries(moduleReexports)) {
+      const targetExports = currentImportMap.get(filename);
+      const moduleExports = currentImportMap.get(name);
+      assertExists(targetExports);
+      assertExists(moduleExports);
+      for (const [exportedName, localName] of reexports) {
+        const exportedValue: any = moduleExports.get(exportedName);
+        assertExists(exportedValue);
+        targetExports.set(localName, exportedValue);
+      }
+    }
   };
-  dynamicImportMap.set(pluginName, dynamicImport);
-  return dynamicImport;
-};
 
-const globalThisMap = new Map<string, any>();
+  const PluginProvider = ({ children }: PropsWithChildren) =>
+    createElement(
+      Provider,
+      {
+        store: rootStore,
+      },
+      children
+    );
 
-export const createOrGetGlobalThis = (
+  const evaluatePluginEntry = (pluginName: string) => {
-  pluginName: string,
-  dynamicImport: (moduleName: string) => Promise<any>
-) => {
-  if (globalThisMap.has(pluginName)) {
     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-    return globalThisMap.get(pluginName)!;
+    const currentImportMap = _pluginNestedImportsMap.get(pluginName)!;
-  }
+    const pluginExports = currentImportMap.get('index.js');
-  const pluginGlobalThis = Object.assign(
+    assertExists(pluginExports);
-    Object.create(null),
+    const entryFunction = pluginExports.get('entry');
-    sharedGlobalThis,
+    const cleanup = entryFunction(<PluginContext>{
-    {
+      register: (part, callback) => {
-      // fixme: vite build output bundle will have this, we should remove it
+        entryLogger.info(`Registering ${pluginName} to ${part}`);
-      process: Object.freeze({
+        if (part === 'headerItem') {
-        env: {
+          rootStore.set(pluginHeaderItemAtom, items => ({
-          NODE_ENV: process.env.NODE_ENV,
+            ...items,
-        },
+            [pluginName]: callback as CallbackMap['headerItem'],
-      }),
+          }));
-      // dynamic import function
-      [dynamicImportKey]: dynamicImport,
-      // UNSAFE: React will read `window` and `document`
-      window: new Proxy(
-        {},
-        {
-          get(_, key) {
-            permissionLogger.debug(`${pluginName} is accessing window`, key);
-            if (sharedGlobalThis[key]) return sharedGlobalThis[key];
-            const result = Reflect.get(window, key);
-            if (typeof result === 'function') {
-              if (result === ShadowRoot) {
-                return result;
-              }
-              return function (...args: any[]) {
-                permissionLogger.debug(
-                  `${pluginName} is calling window`,
-                  key,
-                  args
-                );
-                return result.apply(window, args);
-              };
-            }
-            permissionLogger.debug('window', key, result);
-            return result;
-          },
-        }
-      ),
-      document: new Proxy(
-        {},
-        {
-          get(_, key) {
-            permissionLogger.debug(`${pluginName} is accessing document`, key);
-            if (sharedGlobalThis[key]) return sharedGlobalThis[key];
-            const result = Reflect.get(document, key);
-            if (typeof result === 'function') {
-              return function (...args: any[]) {
-                permissionLogger.debug(
-                  `${pluginName} is calling window`,
-                  key,
-                  args
-                );
-                return result.apply(document, args);
-              };
-            }
-            permissionLogger.debug('document', key, result);
-            return result;
-          },
-        }
-      ),
-      navigator: {
-        userAgent: navigator.userAgent,
-      },
-
-      MouseEvent: globalThis.MouseEvent,
-      KeyboardEvent: globalThis.KeyboardEvent,
-      CustomEvent: globalThis.CustomEvent,
-
-      // copilot uses these
-      Date: globalThis.Date,
-      Math: globalThis.Math,
-      URL: globalThis.URL,
-      URLSearchParams: globalThis.URLSearchParams,
-      Headers: globalThis.Headers,
-      TextEncoder: globalThis.TextEncoder,
-      TextDecoder: globalThis.TextDecoder,
-      Request: globalThis.Request,
-
-      // image-preview uses these
-      Blob: globalThis.Blob,
-      ClipboardItem: globalThis.ClipboardItem,
-
-      // vue uses these
-      Element: globalThis.Element,
-      SVGElement: globalThis.SVGElement,
-
-      // fixme: use our own db api
-      indexedDB: globalThis.indexedDB,
-      IDBRequest: globalThis.IDBRequest,
-      IDBDatabase: globalThis.IDBDatabase,
-      IDBCursorWithValue: globalThis.IDBCursorWithValue,
-      IDBFactory: globalThis.IDBFactory,
-      IDBKeyRange: globalThis.IDBKeyRange,
-      IDBOpenDBRequest: globalThis.IDBOpenDBRequest,
-      IDBTransaction: globalThis.IDBTransaction,
-      IDBObjectStore: globalThis.IDBObjectStore,
-      IDBIndex: globalThis.IDBIndex,
-      IDBCursor: globalThis.IDBCursor,
-      IDBVersionChangeEvent: globalThis.IDBVersionChangeEvent,
-    }
-  );
-  pluginGlobalThis.global = pluginGlobalThis;
-  globalThisMap.set(pluginName, pluginGlobalThis);
-  return pluginGlobalThis;
-};
-
-export const setupPluginCode = async (
-  baseUrl: string,
-  pluginName: string,
-  filename: string
-) => {
-  await rootImportsMapSetupPromise;
-  if (!_pluginNestedImportsMap.has(pluginName)) {
-    _pluginNestedImportsMap.set(pluginName, new Map());
-  }
-  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-  const currentImportMap = _pluginNestedImportsMap.get(pluginName)!;
-  const isMissingPackage = (name: string) =>
-    _rootImportsMap.has(name) && !currentImportMap.has(name);
-
-  const bundleAnalysis = await fetch(`${baseUrl}/${filename}.json`).then(res =>
-    res.json()
-  );
-  const moduleExports = bundleAnalysis.exports as Record<string, [string]>;
-  const moduleImports = bundleAnalysis.imports as string[];
-  const moduleReexports = bundleAnalysis.reexports as Record<
-    string,
-    [localName: string, exportedName: string][]
-  >;
-  await Promise.all(
-    moduleImports.map(name => {
-      if (isMissingPackage(name)) {
-        return Promise.resolve();
-      } else {
-        importLogger.debug('missing package', name);
-        return setupPluginCode(baseUrl, pluginName, name);
-      }
-    })
-  );
-  const code = await fetch(`${baseUrl}/${filename.replace(/^\.\//, '')}`).then(
-    res => res.text()
-  );
-  importLogger.debug('evaluating', filename);
-  const moduleCompartment = new Compartment(
-    createOrGetGlobalThis(
-      pluginName,
-      // use singleton here to avoid infinite loop
-      createOrGetDynamicImport(baseUrl, pluginName)
-    )
-  );
-  const entryPoint = moduleCompartment.evaluate(code, {
-    __evadeHtmlCommentTest__: true,
-  });
-  const moduleExportsMap = new Map<string, any>();
-  const setVarProxy = new Proxy(
-    {},
-    {
-      get(_, p: string): any {
-        return (newValue: any) => {
-          moduleExportsMap.set(p, newValue);
-        };
-      },
-    }
-  );
-  currentImportMap.set(filename, moduleExportsMap);
-  entryPoint({
-    imports: createImports(pluginName),
-    liveVar: setVarProxy,
-    onceVar: setVarProxy,
-  });
-
-  for (const [newExport, [originalExport]] of Object.entries(moduleExports)) {
-    if (newExport === originalExport) continue;
-    const value = moduleExportsMap.get(originalExport);
-    moduleExportsMap.set(newExport, value);
-    moduleExportsMap.delete(originalExport);
-  }
-
-  for (const [name, reexports] of Object.entries(moduleReexports)) {
-    const targetExports = currentImportMap.get(filename);
-    const moduleExports = currentImportMap.get(name);
-    assertExists(targetExports);
-    assertExists(moduleExports);
-    for (const [exportedName, localName] of reexports) {
-      const exportedValue: any = moduleExports.get(exportedName);
-      assertExists(exportedValue);
-      targetExports.set(localName, exportedValue);
-    }
-  }
-};
-
-const PluginProvider = ({ children }: PropsWithChildren) =>
-  createElement(
-    Provider,
-    {
-      store: rootStore,
-    },
-    children
-  );
-
-const entryLogger = new DebugLogger('plugin:entry');
-
-export const evaluatePluginEntry = (pluginName: string) => {
-  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-  const currentImportMap = _pluginNestedImportsMap.get(pluginName)!;
-  const pluginExports = currentImportMap.get('index.js');
-  assertExists(pluginExports);
-  const entryFunction = pluginExports.get('entry');
-  const cleanup = entryFunction(<PluginContext>{
-    register: (part, callback) => {
-      entryLogger.info(`Registering ${pluginName} to ${part}`);
-      if (part === 'headerItem') {
-        rootStore.set(pluginHeaderItemAtom, items => ({
-          ...items,
-          [pluginName]: callback as CallbackMap['headerItem'],
-        }));
-        addCleanup(pluginName, () => {
-          rootStore.set(pluginHeaderItemAtom, items => {
-            const newItems = { ...items };
-            delete newItems[pluginName];
-            return newItems;
-          });
-        });
-      } else if (part === 'editor') {
-        rootStore.set(pluginEditorAtom, items => ({
-          ...items,
-          [pluginName]: callback as CallbackMap['editor'],
-        }));
-        addCleanup(pluginName, () => {
-          rootStore.set(pluginEditorAtom, items => {
-            const newItems = { ...items };
-            delete newItems[pluginName];
-            return newItems;
-          });
-        });
-      } else if (part === 'setting') {
-        rootStore.set(pluginSettingAtom, items => ({
-          ...items,
-          [pluginName]: callback as CallbackMap['setting'],
-        }));
-        addCleanup(pluginName, () => {
-          rootStore.set(pluginSettingAtom, items => {
-            const newItems = { ...items };
-            delete newItems[pluginName];
-            return newItems;
-          });
-        });
-      } else if (part === 'formatBar') {
-        const register = (widget: AffineFormatBarWidget) => {
-          const div = document.createElement('div');
-          const root = widget.root;
-          const cleanup = (callback as CallbackMap['formatBar'])(
-            div,
-            widget.page,
-            () => {
-              return root.selectionManager.value;
-            }
-          );
           addCleanup(pluginName, () => {
-            AffineFormatBarWidget.customElements.delete(register);
+            rootStore.set(pluginHeaderItemAtom, items => {
-            cleanup();
+              const newItems = { ...items };
+              delete newItems[pluginName];
+              return newItems;
+            });
           });
-          return div;
+        } else if (part === 'editor') {
-        };
+          rootStore.set(pluginEditorAtom, items => ({
-        AffineFormatBarWidget.customElements.add(register);
+            ...items,
-      } else {
+            [pluginName]: callback as CallbackMap['editor'],
-        throw new Error(`Unknown part: ${part}`);
+          }));
-      }
+          addCleanup(pluginName, () => {
-    },
+            rootStore.set(pluginEditorAtom, items => {
-    utils: {
+              const newItems = { ...items };
-      PluginProvider,
+              delete newItems[pluginName];
-    },
+              return newItems;
-  });
+            });
-  if (typeof cleanup !== 'function') {
+          });
-    throw new Error('Plugin entry must return a function');
+        } else if (part === 'setting') {
-  }
+          rootStore.set(pluginSettingAtom, items => ({
-  addCleanup(pluginName, cleanup);
+            ...items,
-};
+            [pluginName]: callback as CallbackMap['setting'],
+          }));
+          addCleanup(pluginName, () => {
+            rootStore.set(pluginSettingAtom, items => {
+              const newItems = { ...items };
+              delete newItems[pluginName];
+              return newItems;
+            });
+          });
+        } else if (part === 'formatBar') {
+          const register = (widget: AffineFormatBarWidget) => {
+            const div = document.createElement('div');
+            const root = widget.root;
+            const cleanup = (callback as CallbackMap['formatBar'])(
+              div,
+              widget.page,
+              () => {
+                return root.selectionManager.value;
+              }
+            );
+            addCleanup(pluginName, () => {
+              AffineFormatBarWidget.customElements.delete(register);
+              cleanup();
+            });
+            return div;
+          };
+          AffineFormatBarWidget.customElements.add(register);
+        } else {
+          throw new Error(`Unknown part: ${part}`);
+        }
+      },
+      utils: {
+        PluginProvider,
+      },
+    });
+    if (typeof cleanup !== 'function') {
+      throw new Error('Plugin entry must return a function');
+    }
+    addCleanup(pluginName, cleanup);
+  };
+  return {
+    _rootImportsMap,
+    _pluginNestedImportsMap,
+    createImports,
+    createOrGetDynamicImport,
+    setupPluginCode,
+    evaluatePluginEntry,
+    createOrGetGlobalThis,
+  };
+}

apps/core/src/bootstrap/register-plugins.ts

@@ -5,11 +5,14 @@ import {
   invokeCleanup,
   pluginPackageJson,
 } from '@toeverything/infra/__internal__/plugin';
-import { loadedPluginNameAtom, rootStore } from '@toeverything/infra/atom';
+import {
+  getCurrentStore,
+  loadedPluginNameAtom,
+} from '@toeverything/infra/atom';
 import { packageJsonOutputSchema } from '@toeverything/infra/type';
 import type { z } from 'zod';
 
-import { evaluatePluginEntry, setupPluginCode } from './plugins/setup';
+import { createSetup } from './plugins/setup';
 
 const logger = new DebugLogger('register-plugins');
 
@@ -20,107 +23,112 @@ declare global {
 
 Object.defineProperty(globalThis, '__pluginPackageJson__', {
   get() {
-    return rootStore.get(pluginPackageJson);
+    return getCurrentStore().get(pluginPackageJson);
   },
 });
 
-rootStore.sub(enabledPluginAtom, () => {
+export async function bootstrapPluginSystem(
-  const added = new Set<string>();
+  rootStore: ReturnType<typeof getCurrentStore>
-  const removed = new Set<string>();
+) {
-  const enabledPlugin = new Set(rootStore.get(enabledPluginAtom));
+  const { evaluatePluginEntry, setupPluginCode } = createSetup(rootStore);
-  enabledPlugin.forEach(pluginName => {
+  rootStore.sub(enabledPluginAtom, () => {
-    if (!enabledPluginSet.has(pluginName)) {
+    const added = new Set<string>();
-      added.add(pluginName);
+    const removed = new Set<string>();
-    }
+    const enabledPlugin = new Set(rootStore.get(enabledPluginAtom));
+    enabledPlugin.forEach(pluginName => {
+      if (!enabledPluginSet.has(pluginName)) {
+        added.add(pluginName);
+      }
+    });
+    enabledPluginSet.forEach(pluginName => {
+      if (!enabledPlugin.has(pluginName)) {
+        removed.add(pluginName);
+      }
+    });
+    // update plugins
+    enabledPluginSet.clear();
+    enabledPlugin.forEach(pluginName => {
+      enabledPluginSet.add(pluginName);
+    });
+    added.forEach(pluginName => {
+      evaluatePluginEntry(pluginName);
+    });
+    removed.forEach(pluginName => {
+      invokeCleanup(pluginName);
+    });
   });
-  enabledPluginSet.forEach(pluginName => {
+  const enabledPluginSet = new Set(rootStore.get(enabledPluginAtom));
-    if (!enabledPlugin.has(pluginName)) {
+  const loadedAssets = new Set<string>();
-      removed.add(pluginName);
-    }
-  });
-  // update plugins
-  enabledPluginSet.clear();
-  enabledPlugin.forEach(pluginName => {
-    enabledPluginSet.add(pluginName);
-  });
-  added.forEach(pluginName => {
-    evaluatePluginEntry(pluginName);
-  });
-  removed.forEach(pluginName => {
-    invokeCleanup(pluginName);
-  });
-});
-const enabledPluginSet = new Set(rootStore.get(enabledPluginAtom));
-const loadedAssets = new Set<string>();
 
-// we will load all plugins in parallel from builtinPlugins
+  // we will load all plugins in parallel from builtinPlugins
-export const pluginRegisterPromise = Promise.all(
+  return Promise.all(
-  [...builtinPluginPaths].map(url => {
+    [...builtinPluginPaths].map(url => {
-    return fetch(`${url}/package.json`)
+      return fetch(`${url}/package.json`)
-      .then(async res => {
+        .then(async res => {
-        const packageJson = (await res.json()) as z.infer<
+          const packageJson = (await res.json()) as z.infer<
-          typeof packageJsonOutputSchema
+            typeof packageJsonOutputSchema
-        >;
+          >;
-        packageJsonOutputSchema.parse(packageJson);
+          packageJsonOutputSchema.parse(packageJson);
-        const {
+          const {
-          name: pluginName,
+            name: pluginName,
-          affinePlugin: {
+            affinePlugin: {
-            release,
+              release,
-            entry: { core },
+              entry: { core },
-            assets,
+              assets,
-          },
+            },
-        } = packageJson;
+          } = packageJson;
-        rootStore.set(pluginPackageJson, json => [...json, packageJson]);
+          rootStore.set(pluginPackageJson, json => [...json, packageJson]);
-        logger.debug(`registering plugin ${pluginName}`);
+          logger.debug(`registering plugin ${pluginName}`);
-        logger.debug(`package.json: ${packageJson}`);
+          logger.debug(`package.json: ${packageJson}`);
-        if (!release && !runtimeConfig.enablePlugin) {
+          if (!release && !runtimeConfig.enablePlugin) {
-          return Promise.resolve();
+            return Promise.resolve();
-        }
+          }
-        const baseURL = url;
+          const baseURL = url;
-        const entryURL = `${baseURL}/${core}`;
+          const entryURL = `${baseURL}/${core}`;
-        rootStore.set(loadedPluginNameAtom, prev => [...prev, pluginName]);
+          rootStore.set(loadedPluginNameAtom, prev => [...prev, pluginName]);
-        await setupPluginCode(baseURL, pluginName, core);
+          await setupPluginCode(baseURL, pluginName, core);
-        console.log(`prepareImports for ${pluginName} done`);
+          console.log(`prepareImports for ${pluginName} done`);
-        await fetch(entryURL).then(async () => {
+          await fetch(entryURL).then(async () => {
-          if (assets.length > 0) {
+            if (assets.length > 0) {
-            await Promise.all(
+              await Promise.all(
-              assets.map(async (asset: string) => {
+                assets.map(async (asset: string) => {
-                const loadedAssetName = `${pluginName}_${asset}`;
+                  const loadedAssetName = `${pluginName}_${asset}`;
-                // todo(himself65): add assets into shadow dom
+                  // todo(himself65): add assets into shadow dom
-                if (loadedAssets.has(loadedAssetName)) {
+                  if (loadedAssets.has(loadedAssetName)) {
-                  return Promise.resolve();
+                    return Promise.resolve();
-                }
-                if (asset.endsWith('.css')) {
-                  loadedAssets.add(loadedAssetName);
-                  const res = await fetch(`${baseURL}/${asset}`);
-                  if (res.ok) {
-                    // todo: how to put css file into sandbox?
-                    return res.text().then(text => {
-                      const style = document.createElement('style');
-                      style.setAttribute('plugin-id', pluginName);
-                      style.textContent = text;
-                      document.head.appendChild(style);
-                    });
                   }
-                  return null;
+                  if (asset.endsWith('.css')) {
-                } else {
+                    loadedAssets.add(loadedAssetName);
-                  return Promise.resolve();
+                    const res = await fetch(`${baseURL}/${asset}`);
-                }
+                    if (res.ok) {
-              })
+                      // todo: how to put css file into sandbox?
-            );
+                      return res.text().then(text => {
-          }
+                        const style = document.createElement('style');
-          if (!enabledPluginSet.has(pluginName)) {
+                        style.setAttribute('plugin-id', pluginName);
-            logger.debug(`plugin ${pluginName} is not enabled`);
+                        style.textContent = text;
-          } else {
+                        document.head.appendChild(style);
-            logger.debug(`plugin ${pluginName} is enabled`);
+                      });
-            evaluatePluginEntry(pluginName);
+                    }
-          }
+                    return null;
+                  } else {
+                    return Promise.resolve();
+                  }
+                })
+              );
+            }
+            if (!enabledPluginSet.has(pluginName)) {
+              logger.debug(`plugin ${pluginName} is not enabled`);
+            } else {
+              logger.debug(`plugin ${pluginName} is enabled`);
+              evaluatePluginEntry(pluginName);
+            }
+          });
+        })
+        .catch(e => {
+          console.error(`error when fetch plugin from ${url}`, e);
         });
-      })
+    })
-      .catch(e => {
+  ).then(() => {
-        console.error(`error when fetch plugin from ${url}`, e);
+    console.info('All plugins loaded');
-      });
+  });
-  })
+}
-).then(() => {
-  console.info('All plugins loaded');
-});

apps/core/src/bootstrap/setup.ts

@@ -21,7 +21,7 @@ import {
 } from '@affine/workspace/migration';
 import { createIndexedDBDownloadProvider } from '@affine/workspace/providers';
 import { assertExists } from '@blocksuite/global/utils';
-import { rootStore } from '@toeverything/infra/atom';
+import type { createStore } from 'jotai/vanilla';
 
 import { WorkspaceAdapters } from '../adapters/workspace';
 
@@ -143,7 +143,7 @@ async function tryMigration() {
   }
 }
 
-function createFirstAppData() {
+function createFirstAppData(store: ReturnType<typeof createStore>) {
   const createFirst = (): RootWorkspaceMetadataV2[] => {
     const Plugins = Object.values(WorkspaceAdapters).sort(
       (a, b) => a.loadPriority - b.loadPriority
@@ -166,18 +166,11 @@ function createFirstAppData() {
   const result = createFirst();
   console.info('create first workspace', result);
   localStorage.setItem('is-first-open', 'false');
-  rootStore.set(rootWorkspacesMetadataAtom, result);
+  store.set(rootWorkspacesMetadataAtom, result);
 }
 
-let isSetup = false;
+export async function setup(store: ReturnType<typeof createStore>) {
-
+  store.set(
-export async function setup() {
-  if (isSetup) {
-    console.warn('already setup');
-    return;
-  }
-  isSetup = true;
-  rootStore.set(
     workspaceAdaptersAtom,
     WorkspaceAdapters as Record<
       WorkspaceFlavour,
@@ -188,8 +181,8 @@ export async function setup() {
   console.log('setup global');
   setupGlobal();
 
-  createFirstAppData();
+  createFirstAppData(store);
   await tryMigration();
-  await rootStore.get(rootWorkspacesMetadataAtom);
+  await store.get(rootWorkspacesMetadataAtom);
   console.log('setup done');
 }

apps/core/src/components/affine/affine-error-eoundary.tsx

@@ -8,7 +8,7 @@ import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import {
   currentPageIdAtom,
   currentWorkspaceIdAtom,
-  rootStore,
+  getCurrentStore,
 } from '@toeverything/infra/atom';
 import { useAtomValue } from 'jotai/react';
 import { Provider } from 'jotai/react';
@@ -102,7 +102,7 @@ export class AffineErrorBoundary extends Component<
       return (
         <>
           {errorDetail}
-          <Provider key="JotaiProvider" store={rootStore}>
+          <Provider key="JotaiProvider" store={getCurrentStore()}>
             <DumpInfo />
           </Provider>
         </>

apps/core/src/components/page-detail-editor.tsx

@@ -13,7 +13,7 @@ import {
   pluginEditorAtom,
   pluginWindowAtom,
 } from '@toeverything/infra/__internal__/plugin';
-import { contentLayoutAtom, rootStore } from '@toeverything/infra/atom';
+import { contentLayoutAtom, getCurrentStore } from '@toeverything/infra/atom';
 import clsx from 'clsx';
 import { useAtomValue, useSetAtom } from 'jotai';
 import type { CSSProperties, ReactElement } from 'react';
@@ -103,6 +103,7 @@ const EditorWrapper = memo(function EditorWrapper({
             if (onLoad) {
               dispose = onLoad(page, editor);
             }
+            const rootStore = getCurrentStore();
             const editorItems = rootStore.get(pluginEditorAtom);
             let disposes: (() => void)[] = [];
             const renderTimeout = setTimeout(() => {

apps/core/src/hooks/use-workspaces.ts

@@ -5,7 +5,7 @@ import { saveWorkspaceToLocalStorage } from '@affine/workspace/local/crud';
 import { getOrCreateWorkspace } from '@affine/workspace/manager';
 import { nanoid } from '@blocksuite/store';
 import { getWorkspace } from '@toeverything/infra/__internal__/workspace';
-import { rootStore } from '@toeverything/infra/atom';
+import { getCurrentStore } from '@toeverything/infra/atom';
 import { buildShowcaseWorkspace } from '@toeverything/infra/blocksuite';
 import { useAtomValue, useSetAtom } from 'jotai';
 import { useCallback } from 'react';
@@ -55,7 +55,7 @@ export function useAppHelper() {
             WorkspaceFlavour.LOCAL
           );
           await buildShowcaseWorkspace(blockSuiteWorkspace, {
-            store: rootStore,
+            store: getCurrentStore(),
             atoms: {
               pageMode: setPageModeAtom,
             },

apps/core/src/index.tsx

@@ -1,11 +1,18 @@
 import { WorkspaceFallback } from '@affine/component/workspace';
 import { assertExists } from '@blocksuite/global/utils';
+import { getCurrentStore } from '@toeverything/infra/atom';
 import { StrictMode, Suspense } from 'react';
 import { createRoot } from 'react-dom/client';
 
+import { bootstrapPluginSystem } from './bootstrap/register-plugins';
+
 async function main() {
   const { setup } = await import('./bootstrap/setup');
-  await setup();
+  const rootStore = getCurrentStore();
+  await setup(rootStore);
+  bootstrapPluginSystem(rootStore).catch(err => {
+    console.error('Failed to bootstrap plugin system', err);
+  });
   const { App } = await import('./app');
   const root = document.getElementById('app');
   assertExists(root);

apps/core/src/pages/index.tsx

@@ -2,7 +2,7 @@ import { DebugLogger } from '@affine/debug';
 import { DEFAULT_HELLO_WORLD_PAGE_ID_SUFFIX } from '@affine/env/constant';
 import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import { getWorkspace } from '@toeverything/infra/__internal__/workspace';
-import { rootStore } from '@toeverything/infra/atom';
+import { getCurrentStore } from '@toeverything/infra/atom';
 import { lazy } from 'react';
 import type { LoaderFunction } from 'react-router-dom';
 import { redirect } from 'react-router-dom';
@@ -16,6 +16,7 @@ const AllWorkspaceModals = lazy(() =>
 const logger = new DebugLogger('index-page');
 
 export const loader: LoaderFunction = async () => {
+  const rootStore = getCurrentStore();
   const meta = await rootStore.get(rootWorkspacesMetadataAtom);
   const lastId = localStorage.getItem('last_workspace_id');
   const lastPageId = localStorage.getItem('last_page_id');

apps/core/src/pages/workspace/all-page.tsx

@@ -3,7 +3,7 @@ import { DEFAULT_HELLO_WORLD_PAGE_ID_SUFFIX } from '@affine/env/constant';
 import { WorkspaceSubPath } from '@affine/env/workspace';
 import { assertExists } from '@blocksuite/global/utils';
 import { getActiveBlockSuiteWorkspaceAtom } from '@toeverything/infra/__internal__/workspace';
-import { rootStore } from '@toeverything/infra/atom';
+import { getCurrentStore } from '@toeverything/infra/atom';
 import { useCallback } from 'react';
 import type { LoaderFunction } from 'react-router-dom';
 import { redirect } from 'react-router-dom';
@@ -13,6 +13,7 @@ import { useCurrentWorkspace } from '../../hooks/current/use-current-workspace';
 import { useNavigateHelper } from '../../hooks/use-navigate-helper';
 
 export const loader: LoaderFunction = async args => {
+  const rootStore = getCurrentStore();
   const workspaceId = args.params.workspaceId;
   assertExists(workspaceId);
   const workspaceAtom = getActiveBlockSuiteWorkspaceAtom(workspaceId);

apps/core/src/pages/workspace/detail-page.tsx

@@ -12,7 +12,7 @@ import {
   currentPageIdAtom,
   currentWorkspaceAtom,
   currentWorkspaceIdAtom,
-  rootStore,
+  getCurrentStore,
 } from '@toeverything/infra/atom';
 import { useAtomValue } from 'jotai';
 import { type ReactElement, useCallback } from 'react';
@@ -87,6 +87,7 @@ export const DetailPage = (): ReactElement => {
 };
 
 export const loader: LoaderFunction = async args => {
+  const rootStore = getCurrentStore();
   rootStore.set(contentLayoutAtom, 'editor');
   if (args.params.workspaceId) {
     localStorage.setItem('last_workspace_id', args.params.workspaceId);

apps/core/src/pages/workspace/index.tsx

@@ -2,7 +2,7 @@ import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import {
   currentPageIdAtom,
   currentWorkspaceIdAtom,
-  rootStore,
+  getCurrentStore,
 } from '@toeverything/infra/atom';
 import type { ReactElement } from 'react';
 import { type LoaderFunction, Outlet, redirect } from 'react-router-dom';
@@ -10,6 +10,7 @@ import { type LoaderFunction, Outlet, redirect } from 'react-router-dom';
 import { WorkspaceLayout } from '../../layouts/workspace-layout';
 
 export const loader: LoaderFunction = async args => {
+  const rootStore = getCurrentStore();
   const meta = await rootStore.get(rootWorkspacesMetadataAtom);
   if (!meta.some(({ id }) => id === args.params.workspaceId)) {
     return redirect('/404');

apps/storybook/.storybook/preview.tsx

@@ -5,14 +5,16 @@ import '@toeverything/components/style.css';
 import { createI18n } from '@affine/i18n';
 import { ThemeProvider, useTheme } from 'next-themes';
 import { useDarkMode } from 'storybook-dark-mode';
-import { setup } from '@affine/core/bootstrap/setup';
 import { AffineContext } from '@affine/component/context';
-import { use } from 'foxact/use';
 import useSWR from 'swr';
 import type { Decorator } from '@storybook/react';
+import { createStore } from 'jotai/vanilla';
+import { setup } from '@affine/core/bootstrap/setup';
+import { _setCurrentStore } from '@toeverything/infra/atom';
+import { bootstrapPluginSystem } from '@affine/core/bootstrap/register-plugins';
+import { setupGlobal } from '@affine/env/global';
 
-const setupPromise = setup();
+setupGlobal();
-
 export const parameters = {
   backgrounds: { disable: true },
   actions: { argTypesRegex: '^on[A-Z].*' },
@@ -51,10 +53,23 @@ const ThemeChange = () => {
 };
 
 const withContextDecorator: Decorator = (Story, context) => {
-  use(setupPromise);
+  const { data: store } = useSWR(
+    context.id,
+    async () => {
+      localStorage.clear();
+      const store = createStore();
+      _setCurrentStore(store);
+      await setup(store);
+      await bootstrapPluginSystem(store);
+      return store;
+    },
+    {
+      suspense: true,
+    }
+  );
   return (
     <ThemeProvider>
-      <AffineContext>
+      <AffineContext store={store}>
         <ThemeChange />
         <Story {...context} />
       </AffineContext>

apps/storybook/src/stories/core.stories.tsx

@@ -1,9 +1,7 @@
-import { pluginRegisterPromise } from '@affine/core/bootstrap/register-plugins';
 import { routes } from '@affine/core/router';
 import { assertExists } from '@blocksuite/global/utils';
-import type { Decorator, StoryFn } from '@storybook/react';
+import type { StoryFn } from '@storybook/react';
 import { userEvent, waitFor } from '@storybook/testing-library';
-import { use } from 'foxact/use';
 import { Outlet, useLocation } from 'react-router-dom';
 import {
   reactRouterOutlets,
@@ -11,11 +9,6 @@ import {
   withRouter,
 } from 'storybook-addon-react-router-v6';
 
-const withCleanLocalStorage: Decorator = (Story, context) => {
-  localStorage.clear();
-  return <Story {...context} />;
-};
-
 const FakeApp = () => {
   const location = useLocation();
   // fixme: `key` is a hack to force the storybook to re-render the outlet
@@ -30,10 +23,9 @@ export default {
 };
 
 export const Index: StoryFn = () => {
-  use(pluginRegisterPromise);
   return <FakeApp />;
 };
-Index.decorators = [withRouter, withCleanLocalStorage];
+Index.decorators = [withRouter];
 Index.parameters = {
   reactRouter: reactRouterParameters({
     routing: reactRouterOutlets(routes),
@@ -59,7 +51,7 @@ SettingPage.play = async ({ canvasElement }) => {
   ) as Element;
   await userEvent.click(settingModalBtn);
 };
-SettingPage.decorators = [withRouter, withCleanLocalStorage];
+SettingPage.decorators = [withRouter];
 SettingPage.parameters = {
   reactRouter: reactRouterParameters({
     routing: reactRouterOutlets(routes),
@@ -69,7 +61,7 @@ SettingPage.parameters = {
 export const NotFoundPage: StoryFn = () => {
   return <FakeApp />;
 };
-NotFoundPage.decorators = [withRouter, withCleanLocalStorage];
+NotFoundPage.decorators = [withRouter];
 NotFoundPage.parameters = {
   reactRouter: reactRouterParameters({
     routing: reactRouterOutlets(routes),
@@ -99,7 +91,7 @@ WorkspaceList.play = async ({ canvasElement }) => {
   ) as Element;
   await userEvent.click(currentWorkspace);
 };
-WorkspaceList.decorators = [withRouter, withCleanLocalStorage];
+WorkspaceList.decorators = [withRouter];
 WorkspaceList.parameters = {
   reactRouter: reactRouterParameters({
     routing: reactRouterOutlets(routes),

packages/component/src/components/context/index.tsx

@@ -1,20 +1,24 @@
 import { ProviderComposer } from '@affine/component/provider-composer';
 import { ThemeProvider } from '@affine/component/theme-provider';
-import { rootStore } from '@toeverything/infra/atom';
+import type { createStore } from 'jotai';
 import { Provider } from 'jotai';
 import type { PropsWithChildren } from 'react';
 import { useMemo } from 'react';
 
-export function AffineContext(props: PropsWithChildren) {
+export type AffineContextProps = PropsWithChildren<{
+  store?: ReturnType<typeof createStore>;
+}>;
+
+export function AffineContext(props: AffineContextProps) {
   return (
     <ProviderComposer
       contexts={useMemo(
         () =>
           [
-            <Provider key="JotaiProvider" store={rootStore} />,
+            <Provider key="JotaiProvider" store={props.store} />,
             <ThemeProvider key="ThemeProvider" />,
           ].filter(Boolean),
-        []
+        [props.store]
       )}
     >
       {props.children}

packages/infra/src/atom.ts

@@ -6,7 +6,19 @@ import { atom, createStore } from 'jotai/vanilla';
 import { getWorkspace, waitForWorkspace } from './__internal__/workspace.js';
 
 // global store
-export const rootStore = createStore();
+let rootStore = createStore();
+
+export function getCurrentStore() {
+  return rootStore;
+}
+
+/**
+ * @internal do not use this function unless you know what you are doing
+ */
+export function _setCurrentStore(store: ReturnType<typeof createStore>) {
+  rootStore = store;
+}
+
 export const loadedPluginNameAtom = atom<string[]>([]);
 
 export const currentWorkspaceIdAtom = atom<string | null>(null);