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,8 +111,29 @@ const deleteLayoutAtom = atom<null, [string], void>(null, (_, set, id) => {
   });
 });
 
+const setupWeakMap = new WeakMap<
+  ReturnType<typeof createStore>,
+  ReturnType<typeof createSetupImpl>
+>();
+
+export function createSetup(rootStore: ReturnType<typeof createStore>) {
+  if (setupWeakMap.has(rootStore)) {
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    return setupWeakMap.get(rootStore)!;
+  }
+  const setup = createSetupImpl(rootStore);
+  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[]
-export const _rootImportsMap = new Map<string, Map<string, any>>();
+  const _rootImportsMap = new Map<string, Map<string, any>>();
   const rootImportsMapSetupPromise = setupImportsMap(_rootImportsMap, {
     react: import('react'),
     'react/jsx-runtime': import('react/jsx-runtime'),
@@ -120,7 +146,7 @@ const rootImportsMapSetupPromise = setupImportsMap(_rootImportsMap, {
     '@blocksuite/icons': import('@blocksuite/icons'),
     '@blocksuite/blocks': import('@blocksuite/blocks'),
     '@affine/sdk/entry': {
-    rootStore: rootStore,
+      rootStore,
       currentWorkspaceAtom: currentWorkspaceAtom,
       currentPageAtom: currentPageAtom,
       pushLayoutAtom: pushLayoutAtom,
@@ -128,17 +154,19 @@ const rootImportsMapSetupPromise = setupImportsMap(_rootImportsMap, {
     },
     '@blocksuite/global/utils': import('@blocksuite/global/utils'),
     '@toeverything/infra/atom': import('@toeverything/infra/atom'),
-  '@toeverything/components/button': import('@toeverything/components/button'),
+    '@toeverything/components/button': import(
+      '@toeverything/components/button'
+    ),
   });
 
   // pluginName -> module -> importName -> updater[]
-export const _pluginNestedImportsMap = new Map<
+  const _pluginNestedImportsMap = new Map<
     string,
     Map<string, Map<string, any>>
   >();
 
   const pluginImportsFunctionMap = new Map<string, (imports: any) => void>();
-export const createImports = (pluginName: string) => {
+  const createImports = (pluginName: string) => {
     if (pluginImportsFunctionMap.has(pluginName)) {
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
       return pluginImportsFunctionMap.get(pluginName)!;
@@ -204,10 +232,7 @@ const dynamicImportMap = new Map<
     (moduleName: string) => Promise<any>
   >();
 
-export const createOrGetDynamicImport = (
+  const createOrGetDynamicImport = (baseUrl: string, pluginName: string) => {
-  baseUrl: string,
-  pluginName: string
-) => {
     if (dynamicImportMap.has(pluginName)) {
       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
       return dynamicImportMap.get(pluginName)!;
@@ -255,7 +280,7 @@ export const createOrGetDynamicImport = (
 
   const globalThisMap = new Map<string, any>();
 
-export const createOrGetGlobalThis = (
+  const createOrGetGlobalThis = (
     pluginName: string,
     dynamicImport: (moduleName: string) => Promise<any>
   ) => {
@@ -305,7 +330,10 @@ export const createOrGetGlobalThis = (
           {},
           {
             get(_, key) {
-            permissionLogger.debug(`${pluginName} is accessing document`, key);
+              permissionLogger.debug(
+                `${pluginName} is accessing document`,
+                key
+              );
               if (sharedGlobalThis[key]) return sharedGlobalThis[key];
               const result = Reflect.get(document, key);
               if (typeof result === 'function') {
@@ -369,7 +397,7 @@ export const createOrGetGlobalThis = (
     return pluginGlobalThis;
   };
 
-export const setupPluginCode = async (
+  const setupPluginCode = async (
     baseUrl: string,
     pluginName: string,
     filename: string
@@ -383,8 +411,8 @@ export const setupPluginCode = async (
     const isMissingPackage = (name: string) =>
       _rootImportsMap.has(name) && !currentImportMap.has(name);
 
-  const bundleAnalysis = await fetch(`${baseUrl}/${filename}.json`).then(res =>
+    const bundleAnalysis = await fetch(`${baseUrl}/${filename}.json`).then(
-    res.json()
+      res => res.json()
     );
     const moduleExports = bundleAnalysis.exports as Record<string, [string]>;
     const moduleImports = bundleAnalysis.imports as string[];
@@ -402,9 +430,9 @@ export const setupPluginCode = async (
         }
       })
     );
-  const code = await fetch(`${baseUrl}/${filename.replace(/^\.\//, '')}`).then(
+    const code = await fetch(
-    res => res.text()
+      `${baseUrl}/${filename.replace(/^\.\//, '')}`
-  );
+    ).then(res => res.text());
     importLogger.debug('evaluating', filename);
     const moduleCompartment = new Compartment(
       createOrGetGlobalThis(
@@ -463,9 +491,7 @@ const PluginProvider = ({ children }: PropsWithChildren) =>
       children
     );
 
-const entryLogger = new DebugLogger('plugin:entry');
+  const evaluatePluginEntry = (pluginName: string) => {
-
-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');
@@ -541,3 +567,13 @@ export const evaluatePluginEntry = (pluginName: string) => {
     }
     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,10 +23,14 @@ declare global {
 
 Object.defineProperty(globalThis, '__pluginPackageJson__', {
   get() {
-    return rootStore.get(pluginPackageJson);
+    return getCurrentStore().get(pluginPackageJson);
   },
 });
 
+export async function bootstrapPluginSystem(
+  rootStore: ReturnType<typeof getCurrentStore>
+) {
+  const { evaluatePluginEntry, setupPluginCode } = createSetup(rootStore);
   rootStore.sub(enabledPluginAtom, () => {
     const added = new Set<string>();
     const removed = new Set<string>();
@@ -54,7 +61,7 @@ const enabledPluginSet = new Set(rootStore.get(enabledPluginAtom));
   const loadedAssets = new Set<string>();
 
   // we will load all plugins in parallel from builtinPlugins
-export const pluginRegisterPromise = Promise.all(
+  return Promise.all(
     [...builtinPluginPaths].map(url => {
       return fetch(`${url}/package.json`)
         .then(async res => {
@@ -124,3 +131,4 @@ export const pluginRegisterPromise = Promise.all(
   ).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);