refactor: workspace manager (#5060)

tests/affine-cloud/e2e/collaboration.spec.ts

@@ -151,10 +151,10 @@ test('can collaborate with other user and name should display when editing', asy
   {
     const title = getBlockSuiteEditorTitle(page2);
     expect(await title.innerText()).toBe('TEST TITLE');
-    const typingPromise = Promise.all([
-      page.keyboard.press('Enter', { delay: 50 }),
-      page.keyboard.type('TEST CONTENT', { delay: 50 }),
-    ]);
+    const typingPromise = (async () => {
+      await page.keyboard.press('Enter', { delay: 50 });
+      await page.keyboard.type('TEST CONTENT', { delay: 50 });
+    })();
     // username should be visible when editing
     await expect(page2.getByText(user.name)).toBeVisible();
     await typingPromise;

tests/affine-desktop/e2e/basic.spec.ts

@@ -18,7 +18,7 @@ test('new page', async ({ page, workspace }) => {
     delay: 100,
   });
   await page.waitForSelector('v-line');
-  const flavour = (await workspace.current()).flavour;
+  const flavour = (await workspace.current()).meta.flavour;
   expect(flavour).toBe('local');
 });
 
@@ -169,14 +169,8 @@ test('windows only check', async ({ page }) => {
 test('delete workspace', async ({ page }) => {
   await clickSideBarCurrentWorkspaceBanner(page);
   await page.getByTestId('new-workspace').click();
-  await page
-    .getByTestId('create-workspace-input')
-    .pressSequentially('Delete Me', {
-      delay: 100,
-    });
-  await page.getByTestId('create-workspace-create-button').click({
-    delay: 100,
-  });
+  await page.getByTestId('create-workspace-input').fill('Delete Me');
+  await page.getByTestId('create-workspace-create-button').click();
   // await page.getByTestId('create-workspace-continue-button').click({
   //   delay: 100,
   // });
@@ -197,9 +191,7 @@ test('delete workspace', async ({ page }) => {
   );
   await page.mouse.wheel(0, 500);
   await page.getByTestId('delete-workspace-button').click();
-  await page
-    .getByTestId('delete-workspace-input')
-    .pressSequentially('Delete Me');
+  await page.getByTestId('delete-workspace-input').fill('Delete Me');
   await page.getByTestId('delete-workspace-confirm-button').click();
   await page.waitForTimeout(1000);
   expect(await page.getByTestId('workspace-name').textContent()).toBe(

tests/affine-desktop/e2e/workspace.spec.ts

@@ -10,7 +10,7 @@ test('check workspace has a DB file', async ({ appInfo, workspace }) => {
   const dbPath = path.join(
     appInfo.sessionData,
     'workspaces',
-    w.id,
+    w.meta.id,
     'storage.db'
   );
   // check if db file exists
@@ -25,7 +25,7 @@ test.skip('move workspace db file', async ({ page, appInfo, workspace }) => {
   // goto workspace setting
   await page.getByTestId('workspace-list-item').click();
 
-  const tmpPath = path.join(appInfo.sessionData, w.id + '-tmp-dir');
+  const tmpPath = path.join(appInfo.sessionData, w.meta.id + '-tmp-dir');
 
   // move db file to tmp folder
   await page.evaluate(tmpPath => {
@@ -53,7 +53,7 @@ test.fixme('export then add', async ({ page, appInfo, workspace }) => {
   await page.getByTestId('slider-bar-workspace-setting-button').click();
   await expect(page.getByTestId('setting-modal')).toBeVisible();
 
-  const originalId = w.id;
+  const originalId = w.meta.id;
 
   const newWorkspaceName = 'new-test-name';
 
@@ -67,7 +67,7 @@ test.fixme('export then add', async ({ page, appInfo, workspace }) => {
   await page.getByTestId('save-workspace-name').click();
   await page.waitForSelector('text="Update workspace name success"');
 
-  const tmpPath = path.join(appInfo.sessionData, w.id + '-tmp.db');
+  const tmpPath = path.join(appInfo.sessionData, w.meta.id + '-tmp.db');
 
   // export db file to tmp folder
   await page.evaluate(tmpPath => {
@@ -105,7 +105,7 @@ test.fixme('export then add', async ({ page, appInfo, workspace }) => {
   // sleep for a while to wait for the workspace to be added :D
   await page.waitForTimeout(2000);
   const newWorkspace = await workspace.current();
-  expect(newWorkspace.id).not.toBe(originalId);
+  expect(newWorkspace.meta.id).not.toBe(originalId);
   // check its name is correct
   await expect(page.getByTestId('workspace-name')).toHaveText(newWorkspaceName);
 

tests/affine-local/e2e/local-first-avatar.spec.ts

@@ -57,5 +57,5 @@ test('should create a page with a local first avatar and remove it', async ({
 
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/local-first-delete-page.spec.ts

@@ -41,7 +41,7 @@ test('page delete -> refresh page -> it should be disappear', async ({
 
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('page delete -> create new page -> refresh page -> new page should be appear -> old page should be disappear', async ({
@@ -95,7 +95,7 @@ test('page delete -> create new page -> refresh page -> new page should be appea
 
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('delete multiple pages -> create multiple pages -> refresh', async ({
@@ -155,5 +155,5 @@ test('delete multiple pages -> create multiple pages -> refresh', async ({
 
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/local-first-delete-workspace.spec.ts

@@ -45,7 +45,7 @@ test('Create new workspace, then delete it', async ({ page, workspace }) => {
   );
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('Delete last workspace', async ({ page }) => {

tests/affine-local/e2e/local-first-export-page.spec.ts

@@ -41,7 +41,7 @@ test.skip('New a page ,then open it and export html', async ({
   );
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test.skip('New a page ,then open it and export markdown', async ({
@@ -74,5 +74,5 @@ test.skip('New a page ,then open it and export markdown', async ({
   );
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/local-first-favorite-page.spec.ts

@@ -31,7 +31,7 @@ test('New a page and open it ,then favorite it', async ({
   await favoriteBtn.click();
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('Export to html, markdown and png', async ({ page }) => {
@@ -121,5 +121,5 @@ test('Cancel favorite', async ({ page, workspace }) => {
   ).not.toBeUndefined();
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/local-first-favorites-items.spec.ts

@@ -33,7 +33,7 @@ test('Show favorite items in sidebar', async ({ page, workspace }) => {
   );
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('Show favorite reference in sidebar', async ({ page, workspace }) => {
@@ -73,7 +73,7 @@ test('Show favorite reference in sidebar', async ({ page, workspace }) => {
   ).toBeVisible();
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test("Deleted page's reference will not be shown in sidebar", async ({
@@ -124,7 +124,7 @@ test("Deleted page's reference will not be shown in sidebar", async ({
   expect(collapseButton).toHaveAttribute('data-disabled', 'true');
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('Add new favorite page via sidebar', async ({ page }) => {

tests/affine-local/e2e/local-first-new-page.spec.ts

@@ -16,7 +16,7 @@ test('click btn new page', async ({ page, workspace }) => {
   expect(newPageId).not.toBe(originPageId);
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('click btn bew page and find it in all pages', async ({
@@ -33,5 +33,5 @@ test('click btn bew page and find it in all pages', async ({
   expect(cell).not.toBeUndefined();
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/local-first-openpage-newtab.spec.ts

@@ -28,5 +28,5 @@ test('click btn bew page and open in tab', async ({ page, workspace }) => {
   expect(newTabPage.url()).toBe(newPageUrl);
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/local-first-restore-page.spec.ts

@@ -47,5 +47,5 @@ test('New a page , then delete it in all pages, restore it', async ({
   expect(restoreCell).not.toBeUndefined();
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/local-first-show-delete-modal.spec.ts

@@ -32,7 +32,7 @@ test('New a page ,then open it and show delete modal', async ({
   expect(confirmTip).not.toBeUndefined();
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('New a page ,then go to all pages and show delete modal', async ({
@@ -58,5 +58,5 @@ test('New a page ,then go to all pages and show delete modal', async ({
   expect(confirmTip).not.toBeUndefined();
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/local-first-trash-page.spec.ts

@@ -38,5 +38,5 @@ test('New a page , then delete it in all pages, finally find it in trash', async
   ).not.toBeUndefined();
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/local-first-workspace-list.spec.ts

@@ -25,7 +25,7 @@ test('just one item in the workspace list at first', async ({
   ).not.toBeNull();
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('create one workspace in the workspace list', async ({
@@ -57,7 +57,7 @@ test('create one workspace in the workspace list', async ({
   expect(result1).toBe(11);
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });
 
 test('create multi workspace in the workspace list', async ({
@@ -88,7 +88,7 @@ test('create multi workspace in the workspace list', async ({
 
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 
   await openWorkspaceListModal(page);
   await page.waitForTimeout(1000);
@@ -142,5 +142,5 @@ test('create multi workspace in the workspace list', async ({
 
   const nextWorkspace = await workspace.current();
 
-  expect(currentWorkspace.id).toBe(nextWorkspace.id);
+  expect(currentWorkspace.meta.id).toBe(nextWorkspace.meta.id);
 });

tests/affine-local/e2e/local-first-workspace.spec.ts

@@ -11,5 +11,5 @@ test('preset workspace name', async ({ page, workspace }) => {
   expect(await workspaceName.textContent()).toBe('Demo Workspace');
   const currentWorkspace = await workspace.current();
 
-  expect(currentWorkspace.flavour).toContain('local');
+  expect(currentWorkspace.meta.flavour).toContain('local');
 });

tests/affine-local/e2e/router.spec.ts

@@ -20,6 +20,7 @@ test('goto not found workspace', async ({ page }) => {
   // if doesn't wait for timeout, data won't be saved into indexedDB
   await page.waitForTimeout(1000);
   await page.goto(new URL('/workspace/invalid/all', coreUrl).toString());
-  await page.waitForTimeout(3000);
-  expect(page.url()).toBe(new URL('/404', coreUrl).toString());
+  await expect(page.getByTestId('not-found')).toBeVisible({
+    timeout: 10000,
+  });
 });

tests/affine-migration/e2e/basic.spec.ts

@@ -38,9 +38,6 @@ test('v1 to v4', async ({ page }) => {
   await expect(page.getByTestId('upgrade-workspace-button')).toBeVisible();
   await page.getByTestId('upgrade-workspace-button').click();
 
-  await expect(page.getByText('Refresh Current Page')).toBeVisible();
-  await page.getByTestId('upgrade-workspace-button').click();
-
   await expect(page.getByTestId('page-list-item')).toHaveCount(2);
   await page
     .getByTestId('page-list-item-title-text')
@@ -63,10 +60,6 @@ test('v2 to v4, database migration', async ({ page }) => {
   await expect(page.getByTestId('upgrade-workspace-button')).toBeVisible();
   await page.getByTestId('upgrade-workspace-button').click();
 
-  await expect(page.getByText('Refresh Current Page')).toBeVisible();
-  await page.getByTestId('upgrade-workspace-button').click();
-  await waitForEditorLoad(page);
-
   // check page mode is correct
   await expect(page.locator('v-line').nth(0)).toHaveText('hello');
   await expect(page.locator('affine-database')).toBeVisible();
@@ -84,9 +77,6 @@ test('v3 to v4, surface migration', async ({ page }) => {
 
   await expect(page.getByTestId('upgrade-workspace-button')).toBeVisible();
   await page.getByTestId('upgrade-workspace-button').click();
-
-  await expect(page.getByText('Refresh Current Page')).toBeVisible();
-  await page.getByTestId('upgrade-workspace-button').click();
   await waitForEditorLoad(page);
 
   await page.waitForTimeout(500);
@@ -106,9 +96,6 @@ test('v0 to v4, subdoc migration', async ({ page }) => {
   await expect(page.getByTestId('upgrade-workspace-button')).toBeVisible();
   await page.getByTestId('upgrade-workspace-button').click();
 
-  await expect(page.getByText('Refresh Current Page')).toBeVisible();
-  await page.getByTestId('upgrade-workspace-button').click();
-
   await expect(page.getByTestId('page-list-item')).toHaveCount(2);
   await page
     .getByTestId('page-list-item-title-text')

tests/kit/playwright.ts

@@ -28,8 +28,7 @@ function generateUUID() {
 export const enableCoverage = !!process.env.CI || !!process.env.COVERAGE;
 
 type CurrentWorkspace = {
-  id: string;
-  flavour: string;
+  meta: { id: string; flavour: string };
   blockSuiteWorkspace: Workspace;
 };
 

tests/kit/utils/cloud.ts

@@ -1,6 +1,7 @@
 import { openHomePage } from '@affine-test/kit/utils/load-page';
 import {
   clickNewPageButton,
+  waitForAllPagesLoad,
   waitForEditorLoad,
 } from '@affine-test/kit/utils/page-logic';
 import {
@@ -191,7 +192,7 @@ export async function enableCloudWorkspace(page: Page) {
   await page.getByTestId('confirm-enable-affine-cloud-button').click();
   // wait for upload and delete local workspace
   await page.waitForTimeout(2000);
-  await waitForEditorLoad(page);
+  await waitForAllPagesLoad(page);
   await clickNewPageButton(page);
 }
 

tests/storybook/.storybook/preview.tsx

@@ -9,14 +9,19 @@ import MockSessionContext, {
 import { ThemeProvider, useTheme } from 'next-themes';
 import { useDarkMode } from 'storybook-dark-mode';
 import { AffineContext } from '@affine/component/context';
+import { workspaceManager } from '@affine/workspace';
 import useSWR from 'swr';
 import type { Decorator } from '@storybook/react';
 import { createStore } from 'jotai/vanilla';
 import { _setCurrentStore } from '@toeverything/infra/atom';
-import { setupGlobal } from '@affine/env/global';
+import { setupGlobal, type Environment } from '@affine/env/global';
 
 import type { Preview } from '@storybook/react';
 import { useLayoutEffect, useRef } from 'react';
+import { setup } from '@affine/core/bootstrap/setup';
+import { bootstrapPluginSystem } from '@affine/core/bootstrap/register-plugins';
+import { WorkspaceFlavour } from '@affine/env/workspace';
+import { currentWorkspaceAtom } from '@affine/workspace/atom';
 
 setupGlobal();
 export const parameters = {
@@ -101,37 +106,25 @@ const ThemeChange = () => {
   return null;
 };
 
-const storeMap = new Map<string, ReturnType<typeof createStore>>();
-
-const bootstrapPluginSystemPromise = import(
-  '@affine/core/bootstrap/register-plugins'
-).then(({ bootstrapPluginSystem }) => bootstrapPluginSystem);
-
-const setupPromise = import('@affine/core/bootstrap/setup').then(
-  ({ setup }) => setup
-);
+localStorage.clear();
+const store = createStore();
+_setCurrentStore(store);
+setup();
+bootstrapPluginSystem(store).catch(err => {
+  console.error('Failed to bootstrap plugin system', err);
+});
+workspaceManager
+  .createWorkspace(WorkspaceFlavour.LOCAL, async w => {
+    w.meta.setName('test-workspace');
+  })
+  .then(id => {
+    store.set(
+      currentWorkspaceAtom,
+      workspaceManager.use({ flavour: WorkspaceFlavour.LOCAL, id }).workspace
+    );
+  });
 
 const withContextDecorator: Decorator = (Story, context) => {
-  const { data: store } = useSWR(
-    context.id,
-    async () => {
-      if (storeMap.has(context.id)) {
-        return storeMap.get(context.id);
-      }
-      localStorage.clear();
-      const store = createStore();
-      _setCurrentStore(store);
-      const setup = await setupPromise;
-      await setup(store);
-      const bootstrapPluginSystem = await bootstrapPluginSystemPromise;
-      await bootstrapPluginSystem(store);
-      storeMap.set(context.id, store);
-      return store;
-    },
-    {
-      suspense: true,
-    }
-  );
   return (
     <ThemeProvider>
       <AffineContext store={store}>
@@ -152,14 +145,22 @@ const withPlatformSelectionDecorator: Decorator = (Story, context) => {
     }
     switch (context.globals.platform) {
       case 'desktop-macos':
-        environment.isDesktop = true;
-        environment.isMacOs = true;
-        environment.isWindows = false;
+        environment = {
+          ...environment,
+          isBrowser: true,
+          isDesktop: true,
+          isMacOs: true,
+          isWindows: false,
+        } as Environment;
         break;
       case 'desktop-windows':
-        environment.isDesktop = true;
-        environment.isMacOs = false;
-        environment.isWindows = true;
+        environment = {
+          ...environment,
+          isBrowser: true,
+          isDesktop: true,
+          isMacOs: false,
+          isWindows: true,
+        } as Environment;
         break;
       default:
         globalThis.$AFFINE_SETUP = false;

tests/storybook/package.json

@@ -24,6 +24,7 @@
     "@vitejs/plugin-react": "^4.2.0",
     "concurrently": "^8.2.2",
     "jest-mock": "^29.7.0",
+    "nanoid": "^5.0.3",
     "serve": "^14.2.1",
     "ses": "^0.18.8",
     "storybook": "^7.5.3",

tests/storybook/src/stories/card.stories.tsx

@@ -3,7 +3,6 @@ import { BlockCard } from '@affine/component/card/block-card';
 import { WorkspaceCard } from '@affine/component/card/workspace-card';
 import { Tooltip } from '@affine/component/ui/tooltip';
 import { WorkspaceFlavour } from '@affine/env/workspace';
-import { getOrCreateWorkspace } from '@affine/workspace/manager';
 import {
   EdgelessIcon,
   ExportToHtmlIcon,
@@ -20,13 +19,6 @@ export default {
   },
 } satisfies Meta;
 
-const blockSuiteWorkspace = getOrCreateWorkspace(
-  'blocksuite-local',
-  WorkspaceFlavour.LOCAL
-);
-
-blockSuiteWorkspace.meta.setName('Hello World');
-
 export const AffineWorkspaceCard = () => {
   return (
     <WorkspaceCard

tests/storybook/src/stories/image-preview-modal.stories.tsx

@@ -1,10 +1,12 @@
 import { BlockSuiteEditor } from '@affine/component/block-suite-editor';
-import { WorkspaceFlavour } from '@affine/env/workspace';
 // eslint-disable-next-line @typescript-eslint/no-restricted-imports
 import { ImagePreviewModal } from '@affine/image-preview-plugin/src/component';
-import { getOrCreateWorkspace } from '@affine/workspace/manager';
+import { waitForCurrentWorkspaceAtom } from '@affine/workspace/atom';
+import type { Page } from '@blocksuite/store';
 import type { Meta } from '@storybook/react';
 import { initEmptyPage } from '@toeverything/infra/blocksuite';
+import { useAtomValue } from 'jotai';
+import { useEffect, useState } from 'react';
 import { createPortal } from 'react-dom';
 
 export default {
@@ -12,36 +14,46 @@ export default {
   component: ImagePreviewModal,
 } satisfies Meta;
 
-const workspace = getOrCreateWorkspace('test', WorkspaceFlavour.LOCAL);
-const page = workspace.createPage('page0');
-initEmptyPage(page);
-fetch(new URL('@affine-test/fixtures/large-image.png', import.meta.url))
-  .then(res => res.arrayBuffer())
-  .then(async buffer => {
-    const id = await workspace.blob.set(
-      new Blob([buffer], { type: 'image/png' })
-    );
-    const frameId = page.getBlockByFlavour('affine:note')[0].id;
-    page.addBlock(
-      'affine:paragraph',
-      {
-        text: new page.Text('Please double click the image to preview it.'),
-      },
-      frameId
-    );
-    page.addBlock(
-      'affine:image',
-      {
-        sourceId: id,
-      },
-      frameId
-    );
-  })
-  .catch(err => {
-    console.error('Failed to load large-image.png', err);
-  });
-
 export const Default = () => {
+  const workspace = useAtomValue(waitForCurrentWorkspaceAtom);
+
+  const [page, setPage] = useState<Page | null>(null);
+
+  useEffect(() => {
+    const page = workspace.blockSuiteWorkspace.createPage('page0');
+    initEmptyPage(page);
+    fetch(new URL('@affine-test/fixtures/large-image.png', import.meta.url))
+      .then(res => res.arrayBuffer())
+      .then(async buffer => {
+        const id = await workspace.blockSuiteWorkspace.blob.set(
+          new Blob([buffer], { type: 'image/png' })
+        );
+        const frameId = page.getBlockByFlavour('affine:note')[0].id;
+        page.addBlock(
+          'affine:paragraph',
+          {
+            text: new page.Text('Please double click the image to preview it.'),
+          },
+          frameId
+        );
+        page.addBlock(
+          'affine:image',
+          {
+            sourceId: id,
+          },
+          frameId
+        );
+      })
+      .catch(err => {
+        console.error('Failed to load large-image.png', err);
+      });
+    setPage(page);
+  }, [workspace]);
+
+  if (!page) {
+    return null;
+  }
+
   return (
     <div
       style={{

tests/storybook/src/stories/quick-search/quick-search-main.stories.tsx

@@ -7,11 +7,10 @@ import { CMDKQuickSearchModal } from '@affine/core/components/pure/cmdk';
 import { HighlightLabel } from '@affine/core/components/pure/cmdk/highlight';
 import { WorkspaceFlavour } from '@affine/env/workspace';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
-import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
-import { getOrCreateWorkspace } from '@affine/workspace/manager';
+import { currentWorkspaceAtom } from '@affine/workspace/atom';
 import type { Page } from '@blocksuite/store';
 import type { Meta, StoryFn } from '@storybook/react';
-import { currentWorkspaceIdAtom } from '@toeverything/infra/atom';
+import { useWorkspace } from '@toeverything/hooks/use-workspace';
 import { useStore } from 'jotai';
 import { useEffect, useLayoutEffect, useState } from 'react';
 import { withRouter } from 'storybook-addon-react-router-v6';
@@ -80,19 +79,15 @@ function useRegisterCommands() {
 }
 
 function usePrepareWorkspace() {
+  const workspaceId = 'test-workspace';
   const store = useStore();
+  const workspace = useWorkspace({
+    id: workspaceId,
+    flavour: WorkspaceFlavour.LOCAL,
+  });
   useLayoutEffect(() => {
-    const workspaceId = 'test-workspace';
-    getOrCreateWorkspace(workspaceId, WorkspaceFlavour.LOCAL);
-    store.set(rootWorkspacesMetadataAtom, [
-      {
-        id: workspaceId,
-        flavour: WorkspaceFlavour.LOCAL,
-        version: 4,
-      },
-    ]);
-    store.set(currentWorkspaceIdAtom, workspaceId);
-  }, [store]);
+    store.set(currentWorkspaceAtom, workspace);
+  }, [store, workspace]);
 }
 
 export const CMDKStoryWithCommands: StoryFn = () => {

tests/storybook/src/stories/share-menu.stories.tsx

@@ -1,17 +1,15 @@
 import { toast } from '@affine/component';
 import { PublicLinkDisableModal } from '@affine/component/disable-public-link';
 import { ShareMenu } from '@affine/core/components/affine/share-page-modal/share-menu';
-import type {
-  AffineCloudWorkspace,
-  LocalWorkspace,
-} from '@affine/env/workspace';
 import { WorkspaceFlavour } from '@affine/env/workspace';
-import { getOrCreateWorkspace } from '@affine/workspace/manager';
-import type { Page } from '@blocksuite/store';
+import { waitForCurrentWorkspaceAtom } from '@affine/workspace/atom';
+import { type Page } from '@blocksuite/store';
 import { expect } from '@storybook/jest';
 import type { Meta, StoryFn } from '@storybook/react';
-import { use } from 'foxact/use';
-import { useState } from 'react';
+import { initEmptyPage } from '@toeverything/infra/blocksuite';
+import { useAtomValue } from 'jotai';
+import { nanoid } from 'nanoid';
+import { useEffect, useState } from 'react';
 
 export default {
   title: 'AFFiNE/ShareMenu',
@@ -21,57 +19,30 @@ export default {
   },
 } satisfies Meta;
 
-async function initPage(page: Page) {
-  await page.waitForLoaded();
-  // Add page block and surface block at root level
-  const pageBlockId = page.addBlock('affine:page', {
-    title: new page.Text('Hello, world!'),
-  });
-  page.addBlock('affine:surface', {}, pageBlockId);
-  const frameId = page.addBlock('affine:note', {}, pageBlockId);
-  page.addBlock(
-    'affine:paragraph',
-    {
-      text: new page.Text('This is a paragraph.'),
-    },
-    frameId
-  );
-  page.resetHistory();
-}
-
-const blockSuiteWorkspace = getOrCreateWorkspace(
-  'test-workspace',
-  WorkspaceFlavour.LOCAL
-);
-
-const promise = Promise.all([
-  initPage(blockSuiteWorkspace.createPage({ id: 'page0' })),
-  initPage(blockSuiteWorkspace.createPage({ id: 'page1' })),
-  initPage(blockSuiteWorkspace.createPage({ id: 'page2' })),
-]);
-
-const localWorkspace: LocalWorkspace = {
-  id: 'test-workspace',
-  flavour: WorkspaceFlavour.LOCAL,
-  blockSuiteWorkspace,
-};
-
-const affineWorkspace: AffineCloudWorkspace = {
-  id: 'test-workspace',
-  flavour: WorkspaceFlavour.AFFINE_CLOUD,
-  blockSuiteWorkspace,
-};
-
 async function unimplemented() {
   toast('work in progress');
 }
 
 export const Basic: StoryFn = () => {
-  use(promise);
+  const workspace = useAtomValue(waitForCurrentWorkspaceAtom);
+
+  const [page, setPage] = useState<Page | null>(null);
+
+  useEffect(() => {
+    const page = workspace.blockSuiteWorkspace.createPage(nanoid());
+    initEmptyPage(page);
+
+    setPage(page);
+  }, [workspace]);
+
+  if (!page) {
+    return <div></div>;
+  }
+
   return (
     <ShareMenu
-      currentPage={blockSuiteWorkspace.getPage('page0') as Page}
-      workspace={localWorkspace}
+      currentPage={page}
+      workspaceMetadata={workspace}
       onEnableAffineCloud={unimplemented}
     />
   );
@@ -95,11 +66,28 @@ Basic.play = async ({ canvasElement }) => {
 };
 
 export const AffineBasic: StoryFn = () => {
-  use(promise);
+  const workspace = useAtomValue(waitForCurrentWorkspaceAtom);
+
+  const [page, setPage] = useState<Page | null>(null);
+
+  useEffect(() => {
+    const page = workspace.blockSuiteWorkspace.createPage(nanoid());
+    initEmptyPage(page);
+
+    setPage(page);
+  }, [workspace]);
+
+  if (!page) {
+    return <div></div>;
+  }
+
   return (
     <ShareMenu
-      currentPage={blockSuiteWorkspace.getPage('page0') as Page}
-      workspace={affineWorkspace}
+      currentPage={page}
+      workspaceMetadata={{
+        ...workspace.meta,
+        flavour: WorkspaceFlavour.AFFINE_CLOUD,
+      }}
       onEnableAffineCloud={unimplemented}
     />
   );
@@ -107,7 +95,6 @@ export const AffineBasic: StoryFn = () => {
 
 export const DisableModal: StoryFn = () => {
   const [open, setOpen] = useState(false);
-  use(promise);
   return (
     <>
       <div onClick={() => setOpen(!open)}>Disable Public Link</div>

tests/storybook/src/stories/workspace-list.stories.tsx

@@ -1,10 +1,8 @@
 import type { WorkspaceListProps } from '@affine/component/workspace-list';
 import { WorkspaceList } from '@affine/component/workspace-list';
-import { WorkspaceFlavour } from '@affine/env/workspace';
-import { getOrCreateWorkspace } from '@affine/workspace/manager';
-import { arrayMove } from '@dnd-kit/sortable';
+import { workspaceListAtom } from '@affine/workspace/atom';
 import type { Meta } from '@storybook/react';
-import { useState } from 'react';
+import { useAtomValue } from 'jotai';
 
 export default {
   title: 'AFFiNE/WorkspaceList',
@@ -15,49 +13,14 @@ export default {
 } satisfies Meta<WorkspaceListProps>;
 
 export const Default = () => {
-  const [items, setItems] = useState(() => {
-    const items = [
-      {
-        id: '1',
-        flavour: WorkspaceFlavour.LOCAL,
-        blockSuiteWorkspace: getOrCreateWorkspace('1', WorkspaceFlavour.LOCAL),
-      },
-      {
-        id: '2',
-        flavour: WorkspaceFlavour.LOCAL,
-        blockSuiteWorkspace: getOrCreateWorkspace('2', WorkspaceFlavour.LOCAL),
-      },
-      {
-        id: '3',
-        flavour: WorkspaceFlavour.LOCAL,
-        blockSuiteWorkspace: getOrCreateWorkspace('3', WorkspaceFlavour.LOCAL),
-      },
-    ] satisfies WorkspaceListProps['items'];
-
-    items.forEach(item => {
-      item.blockSuiteWorkspace.meta.setName(item.id);
-    });
-
-    return items;
-  });
+  const list = useAtomValue(workspaceListAtom);
   return (
     <WorkspaceList
       currentWorkspaceId={null}
-      items={items}
+      items={list}
       onClick={() => {}}
       onSettingClick={() => {}}
-      onDragEnd={event => {
-        const { active, over } = event;
-
-        if (active.id !== over?.id) {
-          setItems(items => {
-            const oldIndex = items.findIndex(item => item.id === active.id);
-            const newIndex = items.findIndex(item => item.id === over?.id);
-
-            return arrayMove(items, oldIndex, newIndex);
-          });
-        }
-      }}
+      onDragEnd={_ => {}}
     />
   );
 };

tests/storybook/tsconfig.json

@@ -1,6 +1,6 @@
 {
   "extends": "../../tsconfig.json",
-  "include": ["./src"],
+  "include": ["./src/**/*", "./.storybook/**/*"],
   "compilerOptions": {
     // Workaround for storybook build
     "baseUrl": "../..",
@@ -18,6 +18,9 @@
     {
       "path": "../../packages/common/env"
     },
+    {
+      "path": "../../packages/common/infra"
+    },
     {
       "path": "../../packages/frontend/workspace"
     },