refactor(core): refactor atom to use di (#5831)

To support multiple instances, this PR removes some atoms and implements them using the new DI system.

removed atom

- `pageSettingsAtom`
- `currentPageIdAtom`
- `currentModeAtom`

packages/common/infra/src/blocksuite/index.ts

@@ -1,4 +1,3 @@
-export * from './initialization';
 export {
   migratePages as forceUpgradePages,
   migrateGuidCompatibility,

packages/common/infra/src/blocksuite/initialization/index.ts

@@ -1,114 +0,0 @@
-import type {
-  JobMiddleware,
-  Page,
-  PageMeta,
-  PageSnapshot,
-  Workspace,
-  WorkspaceInfoSnapshot,
-} from '@blocksuite/store';
-import { Job } from '@blocksuite/store';
-import type { createStore, WritableAtom } from 'jotai/vanilla';
-import { Map as YMap } from 'yjs';
-
-import { getLatestVersions } from '../migration/blocksuite';
-import { replaceIdMiddleware } from './middleware';
-
-export function initEmptyPage(page: Page, title?: string) {
-  page.load(() => {
-    const pageBlockId = page.addBlock('affine:page', {
-      title: new page.Text(title ?? ''),
-    });
-    page.addBlock('affine:surface', {}, pageBlockId);
-    const noteBlockId = page.addBlock('affine:note', {}, pageBlockId);
-    page.addBlock('affine:paragraph', {}, noteBlockId);
-  });
-}
-
-/**
- * FIXME: Use exported json data to instead of building data.
- */
-export async function buildShowcaseWorkspace(
-  workspace: Workspace,
-  {
-    store,
-    atoms,
-  }: {
-    atoms: {
-      pageMode: WritableAtom<
-        undefined,
-        [pageId: string, mode: 'page' | 'edgeless'],
-        void
-      >;
-    };
-    store: ReturnType<typeof createStore>;
-  }
-) {
-  const { onboarding } = await import('@affine/templates');
-
-  const info = onboarding['info.json'] as WorkspaceInfoSnapshot;
-
-  const migrationMiddleware: JobMiddleware = ({ slots, workspace }) => {
-    slots.afterImport.on(payload => {
-      if (payload.type === 'page') {
-        workspace.schema.upgradePage(
-          info?.pageVersion ?? 0,
-          {},
-          payload.page.spaceDoc
-        );
-      }
-    });
-  };
-
-  const job = new Job({
-    workspace,
-    middlewares: [replaceIdMiddleware, migrationMiddleware],
-  });
-
-  job.snapshotToWorkspaceInfo(info);
-
-  // for now all onboarding assets are considered served via CDN
-  // hack assets so that every blob exists
-  // @ts-expect-error - rethinking API
-  job._assetsManager.writeToBlob = async () => {};
-
-  const pageSnapshots: PageSnapshot[] = Object.entries(onboarding)
-    .filter(([key]) => {
-      return key.endsWith('snapshot.json');
-    })
-    .map(([_, value]) => value as unknown as PageSnapshot);
-
-  await Promise.all(
-    pageSnapshots.map(snapshot => {
-      return job.snapshotToPage(snapshot);
-    })
-  );
-
-  const newVersions = getLatestVersions(workspace.schema);
-  workspace.doc
-    .getMap('meta')
-    .set('blockVersions', new YMap(Object.entries(newVersions)));
-
-  // todo: find better way to do the following
-  // perhaps put them into middleware?
-  {
-    // the "AFFiNE - not just a note-taking app" page should be set to edgeless mode
-    const edgelessPage1 = (workspace.meta.pages as PageMeta[])?.find(
-      p => p.title === 'AFFiNE - not just a note-taking app'
-    )?.id;
-
-    if (edgelessPage1) {
-      store.set(atoms.pageMode, edgelessPage1, 'edgeless');
-    }
-
-    // should jump to "AFFiNE - not just a note-taking app" by default
-    const defaultPage = (workspace.meta.pages as PageMeta[])?.find(p =>
-      p.title.startsWith('AFFiNE - not just a note-taking app')
-    )?.id;
-
-    if (defaultPage) {
-      workspace.setPageMeta(defaultPage, {
-        jumpOnce: true,
-      });
-    }
-  }
-}

packages/common/infra/src/di/core/collection.ts

@@ -337,7 +337,7 @@ class ServiceCollectionEditor {
    * ```
    */
   addImpl = <
-    Arg1 extends ServiceIdentifier<any>,
+    Arg1 extends ServiceIdentifier<any> | (new (...args: any) => any),
     Arg2 extends Type<Trait> | ServiceFactory<Trait> | Trait,
     Trait = ServiceIdentifierType<Arg1>,
     Deps extends Arg2 extends Type<Trait>

packages/common/infra/src/index.ts

@@ -3,6 +3,7 @@ export * from './atom';
 export * from './blocksuite';
 export * from './command';
 export * from './di';
+export * from './initialization';
 export * from './livedata';
 export * from './page';
 export * from './storage';

packages/common/infra/src/initialization/index.ts

@@ -0,0 +1,117 @@
+import type { WorkspaceFlavour } from '@affine/env/workspace';
+import type {
+  JobMiddleware,
+  Page,
+  PageSnapshot,
+  WorkspaceInfoSnapshot,
+} from '@blocksuite/store';
+import { Job } from '@blocksuite/store';
+import { Map as YMap } from 'yjs';
+
+import { getLatestVersions } from '../blocksuite/migration/blocksuite';
+import { PageRecordList } from '../page';
+import type { WorkspaceManager } from '../workspace';
+import { replaceIdMiddleware } from './middleware';
+
+export function initEmptyPage(page: Page, title?: string) {
+  page.load(() => {
+    const pageBlockId = page.addBlock('affine:page', {
+      title: new page.Text(title ?? ''),
+    });
+    page.addBlock('affine:surface', {}, pageBlockId);
+    const noteBlockId = page.addBlock('affine:note', {}, pageBlockId);
+    page.addBlock('affine:paragraph', {}, noteBlockId);
+  });
+}
+
+/**
+ * FIXME: Use exported json data to instead of building data.
+ */
+export async function buildShowcaseWorkspace(
+  workspaceManager: WorkspaceManager,
+  flavour: WorkspaceFlavour,
+  workspaceName: string
+) {
+  const meta = await workspaceManager.createWorkspace(
+    flavour,
+    async blockSuiteWorkspace => {
+      blockSuiteWorkspace.meta.setName(workspaceName);
+      const { onboarding } = await import('@affine/templates');
+
+      const info = onboarding['info.json'] as WorkspaceInfoSnapshot;
+
+      const migrationMiddleware: JobMiddleware = ({ slots, workspace }) => {
+        slots.afterImport.on(payload => {
+          if (payload.type === 'page') {
+            workspace.schema.upgradePage(
+              info?.pageVersion ?? 0,
+              {},
+              payload.page.spaceDoc
+            );
+          }
+        });
+      };
+
+      const job = new Job({
+        workspace: blockSuiteWorkspace,
+        middlewares: [replaceIdMiddleware, migrationMiddleware],
+      });
+
+      job.snapshotToWorkspaceInfo(info);
+
+      // for now all onboarding assets are considered served via CDN
+      // hack assets so that every blob exists
+      // @ts-expect-error - rethinking API
+      job._assetsManager.writeToBlob = async () => {};
+
+      const pageSnapshots: PageSnapshot[] = Object.entries(onboarding)
+        .filter(([key]) => {
+          return key.endsWith('snapshot.json');
+        })
+        .map(([_, value]) => value as unknown as PageSnapshot);
+
+      await Promise.all(
+        pageSnapshots.map(snapshot => {
+          return job.snapshotToPage(snapshot);
+        })
+      );
+
+      const newVersions = getLatestVersions(blockSuiteWorkspace.schema);
+      blockSuiteWorkspace.doc
+        .getMap('meta')
+        .set('blockVersions', new YMap(Object.entries(newVersions)));
+    }
+  );
+
+  const { workspace, release } = workspaceManager.open(meta);
+
+  await workspace.engine.sync.waitForLoadedRootDoc();
+
+  const pageRecordList = workspace.services.get(PageRecordList);
+
+  // todo: find better way to do the following
+  // perhaps put them into middleware?
+  {
+    // the "AFFiNE - not just a note-taking app" page should be set to edgeless mode
+    const edgelessPage1 = pageRecordList.records.value.find(
+      p => p.title.value === 'AFFiNE - not just a note-taking app'
+    );
+
+    if (edgelessPage1) {
+      edgelessPage1.setMode('edgeless');
+    }
+
+    // should jump to "AFFiNE - not just a note-taking app" by default
+    const defaultPage = pageRecordList.records.value.find(p =>
+      p.title.value.startsWith('AFFiNE - not just a note-taking app')
+    );
+
+    if (defaultPage) {
+      defaultPage.setMeta({
+        jumpOnce: true,
+      });
+    }
+  }
+  release();
+  return meta;
+}

packages/common/infra/src/livedata/react.ts

@@ -3,13 +3,29 @@ import { useSyncExternalStore } from 'react';
 
 import type { LiveData } from './index';
 
+function noopSubscribe() {
+  return () => {};
+}
+
+function noopGetSnapshot() {
+  return null;
+}
+
 /**
  * subscribe LiveData and return the value.
  */
-export function useLiveData<T>(liveData: LiveData<T>): T {
+export function useLiveData<Input extends LiveData<any> | null | undefined>(
+  liveData: Input
+): NonNullable<Input> extends LiveData<infer T>
+  ? Input extends undefined
+    ? T | undefined
+    : Input extends null
+      ? T | null
+      : T
+  : never {
   return useSyncExternalStore(
-    liveData.reactSubscribe,
-    liveData.reactGetSnapshot
+    liveData ? liveData.reactSubscribe : noopSubscribe,
+    liveData ? liveData.reactGetSnapshot : noopGetSnapshot
   );
 }
 

packages/common/infra/src/page/__tests__/page.spec.ts

@@ -1,31 +0,0 @@
-import { WorkspaceFlavour } from '@affine/env/workspace';
-import { describe, expect, test } from 'vitest';
-
-import { configureInfraServices, configureTestingInfraServices } from '../..';
-import { ServiceCollection } from '../../di';
-import { WorkspaceManager } from '../../workspace';
-import { PageListService } from '..';
-
-describe('Page System', () => {
-  test('basic', async () => {
-    const services = new ServiceCollection();
-    configureInfraServices(services);
-    configureTestingInfraServices(services);
-
-    const provider = services.provider();
-    const workspaceManager = provider.get(WorkspaceManager);
-
-    const { workspace } = workspaceManager.open(
-      await workspaceManager.createWorkspace(WorkspaceFlavour.LOCAL)
-    );
-
-    const pageListService = workspace.services.get(PageListService);
-    expect(pageListService.pages.value.length).toBe(0);
-
-    workspace.blockSuiteWorkspace.createPage({
-      id: 'page0',
-    });
-
-    expect(pageListService.pages.value.length).toBe(1);
-  });
-});

packages/common/infra/src/page/context.ts

@@ -1,21 +1,23 @@
-import type { Page as BlockSuitePage, PageMeta } from '@blocksuite/store';
+import type { Page as BlockSuitePage } from '@blocksuite/store';
 
 import { createIdentifier, type ServiceCollection } from '../di';
+import type { PageRecord } from './record';
 import { PageScope } from './service-scope';
 
 export const BlockSuitePageContext = createIdentifier<BlockSuitePage>(
   'BlockSuitePageContext'
 );
 
-export const PageMetaContext = createIdentifier<PageMeta>('PageMetaContext');
+export const PageRecordContext =
+  createIdentifier<PageRecord>('PageRecordContext');
 
 export function configurePageContext(
   services: ServiceCollection,
   blockSuitePage: BlockSuitePage,
-  pageMeta: PageMeta
+  pageRecord: PageRecord
 ) {
   services
     .scope(PageScope)
-    .addImpl(PageMetaContext, pageMeta)
+    .addImpl(PageRecordContext, pageRecord)
     .addImpl(BlockSuitePageContext, blockSuitePage);
 }

packages/common/infra/src/page/index.ts

@@ -1,25 +1,26 @@
-export * from './context';
-export * from './list';
 export * from './manager';
 export * from './page';
+export * from './record';
+export * from './record-list';
 export * from './service-scope';
 
 import { type ServiceCollection, ServiceProvider } from '../di';
 import { CleanupService } from '../lifecycle';
-import { Workspace, WorkspaceScope } from '../workspace';
-import { BlockSuitePageContext, PageMetaContext } from './context';
-import { PageListService } from './list';
+import { Workspace, WorkspaceLocalState, WorkspaceScope } from '../workspace';
+import { BlockSuitePageContext, PageRecordContext } from './context';
 import { PageManager } from './manager';
 import { Page } from './page';
+import { PageRecordList } from './record-list';
 import { PageScope } from './service-scope';
 
 export function configurePageServices(services: ServiceCollection) {
   services
     .scope(WorkspaceScope)
-    .add(PageListService, [Workspace])
-    .add(PageManager, [Workspace, PageListService, ServiceProvider]);
+    .add(PageManager, [Workspace, PageRecordList, ServiceProvider])
+    .add(PageRecordList, [Workspace, WorkspaceLocalState]);
+
   services
     .scope(PageScope)
     .add(CleanupService)
-    .add(Page, [PageMetaContext, BlockSuitePageContext, ServiceProvider]);
+    .add(Page, [PageRecordContext, BlockSuitePageContext, ServiceProvider]);
 }

packages/common/infra/src/page/manager.ts

@@ -1,10 +1,8 @@
-import type { PageMeta } from '@blocksuite/store';
-
 import type { ServiceProvider } from '../di';
 import { ObjectPool } from '../utils/object-pool';
 import type { Workspace } from '../workspace';
+import type { PageRecordList } from '.';
 import { configurePageContext } from './context';
-import type { PageListService } from './list';
 import { Page } from './page';
 import { PageScope } from './service-scope';
 
@@ -13,28 +11,21 @@ export class PageManager {
 
   constructor(
     private readonly workspace: Workspace,
-    private readonly pageList: PageListService,
+    private readonly pageRecordList: PageRecordList,
     private readonly serviceProvider: ServiceProvider
   ) {}
 
-  openByPageId(pageId: string) {
-    const pageMeta = this.pageList.getPageMetaById(pageId);
-    if (!pageMeta) {
-      throw new Error('Page not found');
+  open(pageId: string) {
+    const pageRecord = this.pageRecordList.record(pageId).value;
+    if (!pageRecord) {
+      throw new Error('Page record not found');
     }
-
-    return this.open(pageMeta);
-  }
-
-  open(pageMeta: PageMeta) {
-    const blockSuitePage = this.workspace.blockSuiteWorkspace.getPage(
-      pageMeta.id
-    );
+    const blockSuitePage = this.workspace.blockSuiteWorkspace.getPage(pageId);
     if (!blockSuitePage) {
       throw new Error('Page not found');
     }
 
-    const exists = this.pool.get(pageMeta.id);
+    const exists = this.pool.get(pageId);
     if (exists) {
       return { page: exists.obj, release: exists.release };
     }
@@ -43,7 +34,7 @@ export class PageManager {
       // avoid to modify the original service collection
       .clone();
 
-    configurePageContext(serviceCollection, blockSuitePage, pageMeta);
+    configurePageContext(serviceCollection, blockSuitePage, pageRecord);
 
     const provider = serviceCollection.provider(
       PageScope,
@@ -52,7 +43,7 @@ export class PageManager {
 
     const page = provider.get(Page);
 
-    const { obj, release } = this.pool.put(pageMeta.id, page);
+    const { obj, release } = this.pool.put(pageId, page);
 
     return { page: obj, release };
   }

packages/common/infra/src/page/page.ts

@@ -1,15 +1,28 @@
 import type { Page as BlockSuitePage } from '@blocksuite/store';
-import { type PageMeta } from '@blocksuite/store';
 import type { ServiceProvider } from '@toeverything/infra/di';
 
-export class Page {
-  get id() {
-    return this.meta.id;
-  }
+import type { PageMode, PageRecord } from './record';
 
+export class Page {
   constructor(
-    public readonly meta: PageMeta,
+    public readonly record: PageRecord,
     public readonly blockSuitePage: BlockSuitePage,
     public readonly services: ServiceProvider
   ) {}
+
+  get id() {
+    return this.record.id;
+  }
+
+  readonly mete = this.record.meta;
+  readonly mode = this.record.mode;
+  readonly title = this.record.title;
+
+  setMode(mode: PageMode) {
+    this.record.setMode(mode);
+  }
+
+  toggleMode() {
+    this.record.toggleMode();
+  }
 }

packages/common/infra/src/page/record-list.ts

@@ -1,24 +1,35 @@
-import type { PageMeta } from '@blocksuite/store';
 import { Observable } from 'rxjs';
 
 import { LiveData } from '../livedata';
-import { SyncEngineStep, type Workspace } from '../workspace';
+import {
+  SyncEngineStep,
+  type Workspace,
+  type WorkspaceLocalState,
+} from '../workspace';
+import { PageRecord } from './record';
 
-export class PageListService {
-  constructor(private readonly workspace: Workspace) {}
+export class PageRecordList {
+  constructor(
+    private readonly workspace: Workspace,
+    private readonly localState: WorkspaceLocalState
+  ) {}
 
-  public readonly pages = LiveData.from<PageMeta[]>(
+  public readonly records = LiveData.from<PageRecord[]>(
     new Observable(subscriber => {
-      subscriber.next(
-        Array.from(this.workspace.blockSuiteWorkspace.meta.pageMetas)
-      );
+      const emit = () => {
+        subscriber.next(
+          this.workspace.blockSuiteWorkspace.meta.pageMetas.map(
+            v => new PageRecord(v.id, this.workspace, this.localState)
+          )
+        );
+      };
+
+      emit();
 
       const dispose =
-        this.workspace.blockSuiteWorkspace.meta.pageMetasUpdated.on(() => {
-          subscriber.next(
-            Array.from(this.workspace.blockSuiteWorkspace.meta.pageMetas)
-          );
-        }).dispose;
+        this.workspace.blockSuiteWorkspace.meta.pageMetasUpdated.on(
+          emit
+        ).dispose;
       return () => {
         dispose();
       };
@@ -44,7 +55,7 @@ export class PageListService {
     false
   );
 
-  public getPageMetaById(id: string) {
-    return this.pages.value.find(page => page.id === id);
+  public record(id: string) {
+    return this.records.map(record => record.find(record => record.id === id));
   }
 }

packages/common/infra/src/page/record.ts

@@ -0,0 +1,65 @@
+import type { PageMeta } from '@blocksuite/store';
+import { Observable } from 'rxjs';
+
+import { LiveData } from '../livedata';
+import type { Workspace, WorkspaceLocalState } from '../workspace';
+
+export type PageMode = 'edgeless' | 'page';
+
+export class PageRecord {
+  constructor(
+    public readonly id: string,
+    private readonly workspace: Workspace,
+    private readonly localState: WorkspaceLocalState
+  ) {}
+
+  meta = LiveData.from<PageMeta>(
+    new Observable(subscriber => {
+      const emit = () => {
+        const meta = this.workspace.blockSuiteWorkspace.meta.pageMetas.find(
+          page => page.id === this.id
+        );
+        if (meta === undefined) {
+          return;
+        }
+        subscriber.next(meta);
+      };
+
+      emit();
+
+      const dispose =
+        this.workspace.blockSuiteWorkspace.meta.pageMetasUpdated.on(
+          emit
+        ).dispose;
+      return () => {
+        dispose();
+      };
+    }),
+    {
+      id: this.id,
+      title: '',
+      tags: [],
+      createDate: 0,
+    }
+  );
+
+  setMeta(meta: Partial<PageMeta>): void {
+    this.workspace.blockSuiteWorkspace.setPageMeta(this.id, meta);
+  }
+
+  mode: LiveData<PageMode> = LiveData.from(
+    this.localState.watch<PageMode>(`page:${this.id}:mode`),
+    'page'
+  ).map(mode => (mode === 'edgeless' ? 'edgeless' : 'page'));
+
+  setMode(mode: PageMode) {
+    this.localState.set(`page:${this.id}:mode`, mode);
+  }
+
+  toggleMode() {
+    this.setMode(this.mode.value === 'edgeless' ? 'page' : 'edgeless');
+    return this.mode.value;
+  }
+
+  title = this.meta.map(meta => meta.title);
+}

packages/common/infra/src/workspace/index.ts

@@ -6,13 +6,14 @@ export * from './list';
 export * from './manager';
 export * from './metadata';
 export * from './service-scope';
+export * from './storage';
 export * from './testing';
 export * from './upgrade';
 export * from './workspace';
 
 import { type ServiceCollection, ServiceProvider } from '../di';
 import { CleanupService } from '../lifecycle';
-import { GlobalCache, GlobalState } from '../storage';
+import { GlobalCache, GlobalState, MemoryMemento } from '../storage';
 import {
   BlockSuiteWorkspaceContext,
   RootYDocContext,
@@ -33,6 +34,7 @@ import { WorkspaceFactory } from './factory';
 import { WorkspaceListProvider, WorkspaceListService } from './list';
 import { WorkspaceManager } from './manager';
 import { WorkspaceScope } from './service-scope';
+import { WorkspaceLocalState } from './storage';
 import {
   TestingLocalWorkspaceFactory,
   TestingLocalWorkspaceListProvider,
@@ -83,5 +85,7 @@ export function configureTestingWorkspaceServices(services: ServiceCollection) {
     )
     .override(WorkspaceFactory('local'), TestingLocalWorkspaceFactory, [
       GlobalState,
-    ]);
+    ])
+    .scope(WorkspaceScope)
+    .override(WorkspaceLocalState, MemoryMemento);
 }

packages/common/infra/src/workspace/storage.ts

@@ -0,0 +1,8 @@
+import { createIdentifier } from '../di';
+import type { Memento } from '../storage';
+
+export interface WorkspaceLocalState extends Memento {}
+
+export const WorkspaceLocalState = createIdentifier<WorkspaceLocalState>(
+  'WorkspaceLocalState'
+);

packages/common/infra/src/workspace/workspace.ts

@@ -10,6 +10,8 @@ import { type WorkspaceMetadata } from './metadata';
 import type { WorkspaceUpgradeController } from './upgrade';
 import { type WorkspaceUpgradeStatus } from './upgrade';
 
+export type { Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
+
 const logger = new DebugLogger('affine:workspace');
 
 export type WorkspaceStatus = {