refactor: workspace loading logic (#1966)

packages/env/package.json

@@ -12,7 +12,8 @@
   },
   "exports": {
     ".": "./src/index.ts",
-    "./constant": "./src/constant.ts"
+    "./constant": "./src/constant.ts",
+    "./blocksuite": "./src/blocksuite.ts"
   },
   "peerDependencies": {
     "@blocksuite/global": "0.0.0-20230409084303-221991d4-nightly"

packages/env/src/blocksuite.ts

@@ -0,0 +1,81 @@
+import { DebugLogger } from '@affine/debug';
+import markdown from '@affine/templates/Welcome-to-AFFiNE.md';
+import { ContentParser } from '@blocksuite/blocks/content-parser';
+import type { Page, Workspace } from '@blocksuite/store';
+import { nanoid } from '@blocksuite/store';
+
+declare global {
+  interface Window {
+    lastImportedMarkdown: string;
+  }
+}
+
+const demoTitle = markdown
+  .split('\n')
+  .splice(0, 1)
+  .join('')
+  .replaceAll('#', '')
+  .trim();
+
+const demoText = markdown.split('\n').slice(1).join('\n');
+
+const logger = new DebugLogger('init-page');
+
+export function initPage(page: Page): void {
+  logger.debug('initEmptyPage', page.id);
+  // Add page block and surface block at root level
+  const isFirstPage = page.meta.init === true;
+  if (isFirstPage) {
+    page.workspace.setPageMeta(page.id, {
+      init: false,
+    });
+    _initPageWithDemoMarkdown(page);
+  } else {
+    _initEmptyPage(page);
+  }
+  page.resetHistory();
+}
+
+export function _initEmptyPage(page: Page, title?: string): void {
+  const pageBlockId = page.addBlock('affine:page', {
+    title: new page.Text(title ?? ''),
+  });
+  page.addBlock('affine:surface', {}, null);
+  const frameId = page.addBlock('affine:frame', {}, pageBlockId);
+  page.addBlock('affine:paragraph', {}, frameId);
+}
+
+export function _initPageWithDemoMarkdown(page: Page): void {
+  logger.debug('initPageWithDefaultMarkdown', page.id);
+  const pageBlockId = page.addBlock('affine:page', {
+    title: new page.Text(demoTitle),
+  });
+  page.addBlock('affine:surface', {}, null);
+  const frameId = page.addBlock('affine:frame', {}, pageBlockId);
+  page.addBlock('affine:paragraph', {}, frameId);
+  const contentParser = new ContentParser(page);
+  contentParser.importMarkdown(demoText, frameId);
+  page.workspace.setPageMeta(page.id, { demoTitle });
+}
+
+export function ensureRootPinboard(blockSuiteWorkspace: Workspace) {
+  const metas = blockSuiteWorkspace.meta.pageMetas;
+  const rootMeta = metas.find(m => m.isRootPinboard);
+
+  if (rootMeta) {
+    return rootMeta.id;
+  }
+
+  const rootPinboardPage = blockSuiteWorkspace.createPage(nanoid());
+
+  const title = `${blockSuiteWorkspace.meta.name}'s Pinboard`;
+
+  _initEmptyPage(rootPinboardPage, title);
+
+  blockSuiteWorkspace.meta.setPageMeta(rootPinboardPage.id, {
+    isRootPinboard: true,
+    title,
+  });
+
+  return rootPinboardPage.id;
+}

packages/env/src/constant.ts

@@ -1,5 +1,6 @@
 export const DEFAULT_WORKSPACE_NAME = 'Demo Workspace';
 export const UNTITLED_WORKSPACE_NAME = 'Untitled';
+export const DEFAULT_HELLO_WORLD_PAGE_ID = 'hello-world';
 
 export const enum MessageCode {
   loginError,

packages/env/src/types.d.ts

@@ -0,0 +1,10 @@
+/// <reference types="@webpack/env"" />
+
+// not using import because it will break the declare module line below
+// eslint-disable-next-line @typescript-eslint/triple-slash-reference
+/// <reference path='../../../electron/layers/preload/preload.d.ts' />
+
+declare module '*.md' {
+  const text: string;
+  export default text;
+}

packages/hooks/src/use-blocksuite-workspace-page-is-public.ts

@@ -10,9 +10,12 @@ declare module '@blocksuite/store' {
 export function useBlockSuiteWorkspacePageIsPublic(page: Page) {
   const [isPublic, set] = useState<boolean>(() => page.meta.isPublic ?? false);
   useEffect(() => {
-    page.workspace.meta.pageMetasUpdated.on(() => {
+    const disposable = page.workspace.meta.pageMetasUpdated.on(() => {
       set(page.meta.isPublic ?? false);
     });
+    return () => {
+      disposable.dispose();
+    };
   }, [page]);
   const setIsPublic = useCallback(
     (isPublic: boolean) => {

packages/hooks/src/use-blocksuite-workspace-page.ts

@@ -0,0 +1,30 @@
+import type { Page, Workspace } from '@blocksuite/store';
+import { useEffect, useState } from 'react';
+
+export function useBlockSuiteWorkspacePage(
+  blockSuiteWorkspace: Workspace,
+  pageId: string | null
+): Page | null {
+  const [page, setPage] = useState(() => {
+    if (pageId === null) {
+      return null;
+    }
+    return blockSuiteWorkspace.getPage(pageId);
+  });
+  useEffect(() => {
+    if (pageId) {
+      setPage(blockSuiteWorkspace.getPage(pageId));
+    }
+  }, [blockSuiteWorkspace, pageId]);
+  useEffect(() => {
+    const disposable = blockSuiteWorkspace.slots.pageAdded.on(id => {
+      if (pageId === id) {
+        setPage(blockSuiteWorkspace.getPage(id));
+      }
+    });
+    return () => {
+      disposable.dispose();
+    };
+  }, [blockSuiteWorkspace, pageId]);
+  return page;
+}

packages/workspace/src/affine/sync.ts

@@ -6,7 +6,7 @@ import {
 } from '@affine/workspace/affine/api';
 import { WebsocketClient } from '@affine/workspace/affine/channel';
 import { storageChangeSlot } from '@affine/workspace/affine/login';
-import { jotaiStore, jotaiWorkspacesAtom } from '@affine/workspace/atom';
+import { rootStore, rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import type { WorkspaceCRUD } from '@affine/workspace/type';
 import type { WorkspaceFlavour } from '@affine/workspace/type';
 import { assertExists } from '@blocksuite/global/utils';
@@ -51,7 +51,7 @@ export function createAffineGlobalChannel(
 
       // If the workspace is not in the current workspace list, remove it
       if (workspaceIndex === -1) {
-        jotaiStore.set(jotaiWorkspacesAtom, workspaces => {
+        rootStore.set(rootWorkspacesMetadataAtom, workspaces => {
           const idx = workspaces.findIndex(workspace => workspace.id === id);
           workspaces.splice(idx, 1);
           return [...workspaces];

packages/workspace/src/atom.ts

@@ -1,18 +1,49 @@
 import { atomWithSyncStorage } from '@affine/jotai';
 import type { WorkspaceFlavour } from '@affine/workspace/type';
-import { createStore } from 'jotai/index';
+import type { EditorContainer } from '@blocksuite/editor';
+import { atom, createStore } from 'jotai';
+import { atomWithStorage, createJSONStorage } from 'jotai/utils';
 
-export type JotaiWorkspace = {
+export type RootWorkspaceMetadata = {
   id: string;
   flavour: WorkspaceFlavour;
 };
+// #region root atoms
+// root primitive atom that stores the necessary data for the whole app
+// be careful when you use this atom,
+// it should be used only in the root component
 
-// root primitive atom that stores the list of workspaces which could be used in the app
-// if a workspace is not in this list, it should not be used in the app
-export const jotaiWorkspacesAtom = atomWithSyncStorage<JotaiWorkspace[]>(
+/**
+ * root workspaces atom
+ * this atom stores the metadata of all workspaces,
+ * which is `id` and `flavour`, that is enough to load the real workspace data
+ */
+export const rootWorkspacesMetadataAtom = atomWithSyncStorage<
+  RootWorkspaceMetadata[]
+>(
+  // don't change this key,
+  // otherwise it will cause the data loss in the production
   'jotai-workspaces',
   []
 );
 
-// global jotai store, which is used to store all the atoms
-export const jotaiStore = createStore();
+// two more atoms to store the current workspace and page
+export const rootCurrentWorkspaceIdAtom = atomWithStorage<string | null>(
+  'root-current-workspace-id',
+  null,
+  createJSONStorage(() => sessionStorage)
+);
+export const rootCurrentPageIdAtom = atomWithStorage<string | null>(
+  'root-current-page-id',
+  null,
+  createJSONStorage(() => sessionStorage)
+);
+
+// current editor atom, each app should have only one editor in the same time
+export const rootCurrentEditorAtom = atom<Readonly<EditorContainer> | null>(
+  null
+);
+//#endregion
+
+// global store
+export const rootStore = createStore();

packages/workspace/src/local/crud.ts

@@ -1,3 +1,4 @@
+import { DebugLogger } from '@affine/debug';
 import { nanoid, Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
 import { createIndexedDBProvider } from '@toeverything/y-indexeddb';
 import { createJSONStorage } from 'jotai/utils';
@@ -13,8 +14,25 @@ const getStorage = () => createJSONStorage(() => localStorage);
 const kStoreKey = 'affine-local-workspace';
 const schema = z.array(z.string());
 
+const logger = new DebugLogger('affine:workspace:local:crud');
+
+/**
+ * @internal
+ */
+export function saveWorkspaceToLocalStorage(workspaceId: string) {
+  const storage = getStorage();
+  !Array.isArray(storage.getItem(kStoreKey)) && storage.setItem(kStoreKey, []);
+  const data = storage.getItem(kStoreKey) as z.infer<typeof schema>;
+  const id = data.find(id => id === workspaceId);
+  if (!id) {
+    logger.debug('saveWorkspaceToLocalStorage', workspaceId);
+    storage.setItem(kStoreKey, [...data, workspaceId]);
+  }
+}
+
 export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
   get: async workspaceId => {
+    logger.debug('get', workspaceId);
     const storage = getStorage();
     !Array.isArray(storage.getItem(kStoreKey)) &&
       storage.setItem(kStoreKey, []);
@@ -36,10 +54,10 @@ export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
     return workspace;
   },
   create: async ({ doc }) => {
+    logger.debug('create', doc);
     const storage = getStorage();
     !Array.isArray(storage.getItem(kStoreKey)) &&
       storage.setItem(kStoreKey, []);
-    const data = storage.getItem(kStoreKey) as z.infer<typeof schema>;
     const binary = BlockSuiteWorkspace.Y.encodeStateAsUpdateV2(doc);
     const id = nanoid();
     const blockSuiteWorkspace = createEmptyBlockSuiteWorkspace(
@@ -52,11 +70,11 @@ export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
     await persistence.whenSynced.then(() => {
       persistence.disconnect();
     });
-    storage.setItem(kStoreKey, [...data, id]);
-    console.log('create', id, storage.getItem(kStoreKey));
+    saveWorkspaceToLocalStorage(id);
     return id;
   },
   delete: async workspace => {
+    logger.debug('delete', workspace);
     const storage = getStorage();
     !Array.isArray(storage.getItem(kStoreKey)) &&
       storage.setItem(kStoreKey, []);
@@ -69,6 +87,7 @@ export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
     storage.setItem(kStoreKey, [...data]);
   },
   list: async () => {
+    logger.debug('list');
     const storage = getStorage();
     !Array.isArray(storage.getItem(kStoreKey)) &&
       storage.setItem(kStoreKey, []);

packages/workspace/src/providers/index.ts

@@ -69,6 +69,36 @@ const createAffineWebSocketProvider = (
   return apis;
 };
 
+class CallbackSet extends Set<() => void> {
+  #ready = false;
+
+  get ready(): boolean {
+    return this.#ready;
+  }
+
+  set ready(v: boolean) {
+    this.#ready = v;
+  }
+
+  add(cb: () => void) {
+    if (this.ready) {
+      cb();
+      return this;
+    }
+    if (this.has(cb)) {
+      return this;
+    }
+    return super.add(cb);
+  }
+
+  delete(cb: () => void) {
+    if (this.has(cb)) {
+      return super.delete(cb);
+    }
+    return false;
+  }
+}
+
 const createIndexedDBProvider = (
   blockSuiteWorkspace: BlockSuiteWorkspace
 ): LocalIndexedDBProvider => {
@@ -76,7 +106,7 @@ const createIndexedDBProvider = (
     blockSuiteWorkspace.id,
     blockSuiteWorkspace.doc
   );
-  const callbacks = new Set<() => void>();
+  const callbacks = new CallbackSet();
   return {
     flavour: 'local-indexeddb',
     callbacks,
@@ -93,6 +123,7 @@ const createIndexedDBProvider = (
       indexeddbProvider.connect();
       indexeddbProvider.whenSynced
         .then(() => {
+          callbacks.ready = true;
           callbacks.forEach(cb => cb());
         })
         .catch(error => {
@@ -110,6 +141,7 @@ const createIndexedDBProvider = (
         blockSuiteWorkspace.id
       );
       indexeddbProvider.disconnect();
+      callbacks.ready = false;
     },
   };
 };

packages/workspace/src/type.ts

@@ -2,8 +2,11 @@
 /// <reference path='../../../apps/electron/layers/preload/preload.d.ts' />
 import type { Workspace as RemoteWorkspace } from '@affine/workspace/affine/api';
 import type { Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
+import type { createStore } from 'jotai';
 import type { FC, PropsWithChildren } from 'react';
 
+export type JotaiStore = ReturnType<typeof createStore>;
+
 export type BaseProvider = {
   flavour: string;
   // if this is true, we will connect the provider on the background
@@ -141,9 +144,9 @@ export interface WorkspaceUISchema<Flavour extends keyof WorkspaceRegistry> {
 }
 
 export interface AppEvents {
-  // event when app is first initialized
+  // event there is no workspace
   // usually used to initialize workspace plugin
-  'app:first-init': () => Promise<void>;
+  'app:init': () => string[];
   // request to gain access to workspace plugin
   'workspace:access': () => Promise<void>;
   // request to revoke access to workspace plugin