feat: run app in closure (#3790)

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);