fix(core): refresh metadata after refresh (#4054)

.eslintrc.js

@@ -37,6 +37,11 @@ const createPattern = packageName => [
     // useSession is type unsafe
     importNames: ['useSession'],
   },
+  {
+    group: ['next-auth/react'],
+    message: "Import hooks from 'cloud-utils.ts'",
+    importNames: ['signIn', 'signOut'],
+  },
   {
     group: ['yjs'],
     message: 'Do not use this API because it has a bug',
@@ -172,6 +177,11 @@ const config = {
             // useSession is type unsafe
             importNames: ['useSession'],
           },
+          {
+            group: ['next-auth/react'],
+            message: "Import hooks from 'cloud-utils.ts'",
+            importNames: ['signIn', 'signOut'],
+          },
           {
             group: ['yjs'],
             message: 'Do not use this API because it has a bug',

apps/core/src/components/affine/auth/after-sign-in-send-email.tsx

@@ -6,9 +6,9 @@ import {
 } from '@affine/component/auth-components';
 import { Trans } from '@affine/i18n';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
-import { signIn } from 'next-auth/react';
 import { type FC, useCallback } from 'react';
 
+import { signInCloud } from '../../../utils/cloud-utils';
 import { buildCallbackUrl } from './callback-url';
 import type { AuthPanelProps } from './index';
 import * as style from './style.css';
@@ -33,7 +33,7 @@ export const AfterSignInSendEmail: FC<AuthPanelProps> = ({
 
       <ResendButton
         onClick={useCallback(() => {
-          signIn('email', {
+          signInCloud('email', {
             email,
             callbackUrl: buildCallbackUrl('signIn'),
             redirect: true,

apps/core/src/components/affine/auth/after-sign-up-send-email.tsx

@@ -5,9 +5,9 @@ import {
   ResendButton,
 } from '@affine/component/auth-components';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
-import { signIn } from 'next-auth/react';
 import { type FC, useCallback } from 'react';
 
+import { signInCloud } from '../../../utils/cloud-utils';
 import { buildCallbackUrl } from './callback-url';
 import type { AuthPanelProps } from './index';
 import * as style from './style.css';
@@ -32,7 +32,7 @@ export const AfterSignUpSendEmail: FC<AuthPanelProps> = ({
 
       <ResendButton
         onClick={useCallback(() => {
-          signIn('email', {
+          signInCloud('email', {
             email: email,
             callbackUrl: buildCallbackUrl('signUp'),
             redirect: true,

apps/core/src/components/affine/auth/index.tsx

@@ -2,8 +2,15 @@ import {
   AuthModal as AuthModalBase,
   type AuthModalProps as AuthModalBaseProps,
 } from '@affine/component/auth-components';
-import { atom, useAtom } from 'jotai';
-import { type FC, useCallback, useEffect, useMemo } from 'react';
+import { refreshRootMetadataAtom } from '@affine/workspace/atom';
+import { atom, useAtom, useSetAtom } from 'jotai';
+import {
+  type FC,
+  startTransition,
+  useCallback,
+  useEffect,
+  useMemo,
+} from 'react';
 
 import { AfterSignInSendEmail } from './after-sign-in-send-email';
 import { AfterSignUpSendEmail } from './after-sign-up-send-email';
@@ -79,9 +86,14 @@ export const AuthModal: FC<AuthModalBaseProps & AuthProps> = ({
     }
   }, [open, setAuthEmail, setAuthStore]);
 
+  const refreshMetadata = useSetAtom(refreshRootMetadataAtom);
+
   const onSignedIn = useCallback(() => {
     setOpen(false);
-  }, [setOpen]);
+    startTransition(() => {
+      refreshMetadata();
+    });
+  }, [refreshMetadata, setOpen]);
 
   return (
     <AuthModalBase open={open} setOpen={setOpen}>

apps/core/src/components/affine/auth/sign-in.tsx

@@ -8,10 +8,11 @@ import { useMutation } from '@affine/workspace/affine/gql';
 import { ArrowDownBigIcon, GoogleDuotoneIcon } from '@blocksuite/icons';
 import { Button } from '@toeverything/components/button';
 import { useSetAtom } from 'jotai';
-import { signIn, type SignInResponse } from 'next-auth/react';
+import { type SignInResponse } from 'next-auth/react';
 import { type FC, useState } from 'react';
 import { useCallback } from 'react';
 
+import { signInCloud } from '../../../utils/cloud-utils';
 import { emailRegex } from '../../../utils/email-regex';
 import { buildCallbackUrl } from './callback-url';
 import type { AuthPanelProps } from './index';
@@ -66,7 +67,7 @@ export const SignIn: FC<AuthPanelProps> = ({
 
     setAuthEmail(email);
     if (user) {
-      signIn('email', {
+      signInCloud('email', {
         email: email,
         callbackUrl: buildCallbackUrl('signIn'),
         redirect: false,
@@ -75,7 +76,7 @@ export const SignIn: FC<AuthPanelProps> = ({
         .catch(console.error);
       setAuthState('afterSignInSendEmail');
     } else {
-      signIn('email', {
+      signInCloud('email', {
         email: email,
         callbackUrl: buildCallbackUrl('signUp'),
         redirect: false,
@@ -102,7 +103,7 @@ export const SignIn: FC<AuthPanelProps> = ({
         }}
         icon={<GoogleDuotoneIcon />}
         onClick={useCallback(() => {
-          signIn('google').catch(console.error);
+          signInCloud('google').catch(console.error);
         }, [])}
       >
         {t['Continue with Google']()}

apps/core/src/components/affine/setting-modal/account-setting/index.tsx

@@ -10,11 +10,11 @@ import { useMutation } from '@affine/workspace/affine/gql';
 import { ArrowRightSmallIcon, CameraIcon, DoneIcon } from '@blocksuite/icons';
 import { Button, IconButton } from '@toeverything/components/button';
 import { useAtom } from 'jotai/index';
-import { signOut } from 'next-auth/react';
 import { type FC, useCallback, useState } from 'react';
 
 import { authAtom } from '../../../../atoms';
 import { useCurrentUser } from '../../../../hooks/affine/use-current-user';
+import { signOutCloud } from '../../../../utils/cloud-utils';
 import { Upload } from '../../../pure/file-upload';
 import * as style from './style.css';
 
@@ -160,7 +160,7 @@ export const AccountSetting: FC = () => {
         desc={t['com.affine.setting.sign.out.message']()}
         style={{ cursor: 'pointer' }}
         onClick={useCallback(() => {
-          signOut().catch(console.error);
+          signOutCloud().catch(console.error);
         }, [])}
       >
         <ArrowRightSmallIcon />

apps/core/src/components/cloud/login-card.tsx

@@ -1,10 +1,10 @@
 import { UserAvatar } from '@affine/component/user-avatar';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import { CloudWorkspaceIcon } from '@blocksuite/icons';
-import { signIn } from 'next-auth/react';
 
 import { useCurrentLoginStatus } from '../../hooks/affine/use-current-login-status';
 import { useCurrentUser } from '../../hooks/affine/use-current-user';
+import { signInCloud } from '../../utils/cloud-utils';
 import { StyledSignInButton } from '../pure/footer/styles';
 
 export const LoginCard = () => {
@@ -17,8 +17,7 @@ export const LoginCard = () => {
     <StyledSignInButton
       data-testid="sign-in-button"
       onClick={async () => {
-        // jump to login page
-        signIn().catch(console.error);
+        signInCloud().catch(console.error);
       }}
     >
       <div className="circle">

apps/core/src/components/pure/footer/index.tsx

@@ -1,12 +1,12 @@
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import { CloudWorkspaceIcon } from '@blocksuite/icons';
-import { signIn } from 'next-auth/react';
 import { type CSSProperties, type FC, forwardRef, useCallback } from 'react';
 
 import { useCurrentLoginStatus } from '../../../hooks/affine/use-current-login-status';
-// import { openDisableCloudAlertModalAtom } from '../../../atoms';
 import { stringToColour } from '../../../utils';
+import { signInCloud } from '../../../utils/cloud-utils';
 import { StyledFooter, StyledSignInButton } from './styles';
+
 export const Footer: FC = () => {
   const loginStatus = useCurrentLoginStatus();
 
@@ -25,7 +25,7 @@ const SignInButton = () => {
     <StyledSignInButton
       data-testid="sign-in-button"
       onClick={useCallback(() => {
-        signIn().catch(console.error);
+        signInCloud().catch(console.error);
       }, [])}
     >
       <div className="circle">

apps/core/src/utils/cloud-utils.tsx

@@ -0,0 +1,24 @@
+import { refreshRootMetadataAtom } from '@affine/workspace/atom';
+import { getCurrentStore } from '@toeverything/infra/atom';
+// eslint-disable-next-line @typescript-eslint/no-restricted-imports
+import { signIn, signOut } from 'next-auth/react';
+import { startTransition } from 'react';
+
+export const signInCloud: typeof signIn = async (...args) => {
+  return signIn(...args).then(result => {
+    // do not refresh root metadata,
+    // because the session won't change in this callback
+    return result;
+  });
+};
+
+export const signOutCloud: typeof signOut = async (...args) => {
+  return signOut(...args).then(result => {
+    if (result) {
+      startTransition(() => {
+        getCurrentStore().set(refreshRootMetadataAtom);
+      });
+    }
+    return result;
+  });
+};

packages/workspace/src/atom.ts

@@ -7,7 +7,7 @@ import {
   currentWorkspaceIdAtom,
 } from '@toeverything/infra/atom';
 import { WorkspaceVersion } from '@toeverything/infra/blocksuite';
-import { atom } from 'jotai';
+import { type Atom, atom } from 'jotai/vanilla';
 import { z } from 'zod';
 
 import { getOrCreateWorkspace } from './manager';
@@ -68,127 +68,145 @@ const METADATA_STORAGE_KEY = 'jotai-workspaces';
 const rootWorkspacesMetadataPrimitiveAtom = atom<Promise<
   RootWorkspaceMetadata[]
 > | null>(null);
+
+type Getter = <Value>(atom: Atom<Value>) => Value;
+
+type FetchMetadata = (
+  get: Getter,
+  options: { signal: AbortSignal }
+) => Promise<RootWorkspaceMetadata[]>;
+
+/**
+ * @internal
+ */
+const fetchMetadata: FetchMetadata = async (get, { signal }) => {
+  const WorkspaceAdapters = get(workspaceAdaptersAtom);
+  assertExists(WorkspaceAdapters, 'workspace adapter should be defined');
+  const metadata: RootWorkspaceMetadata[] = [];
+
+  // step 1: try load metadata from localStorage.
+  //
+  // we need this step because workspaces have the order.
+  {
+    const loadFromLocalStorage = (): RootWorkspaceMetadata[] => {
+      // don't change this key,
+      // otherwise it will cause the data loss in the production
+      const primitiveMetadata = localStorage.getItem(METADATA_STORAGE_KEY);
+      if (primitiveMetadata) {
+        try {
+          const items = JSON.parse(primitiveMetadata) as z.infer<
+            typeof rootWorkspaceMetadataArraySchema
+          >;
+          rootWorkspaceMetadataArraySchema.parse(items);
+          return [...items];
+        } catch (e) {
+          console.error('cannot parse worksapce', e);
+        }
+        return [];
+      }
+      return [];
+    };
+
+    const maybeMetadata = loadFromLocalStorage();
+
+    // migration step, only data in `METADATA_STORAGE_KEY` will be migrated
+    if (
+      maybeMetadata.some(meta => !('version' in meta)) &&
+      !globalThis.$migrationDone
+    ) {
+      await new Promise<void>((resolve, reject) => {
+        signal.addEventListener('abort', () => reject(), { once: true });
+        window.addEventListener('migration-done', () => resolve(), {
+          once: true,
+        });
+      });
+    }
+
+    metadata.push(...loadFromLocalStorage());
+  }
+  // step 2: fetch from adapters
+  {
+    const Adapters = Object.values(WorkspaceAdapters).sort(
+      (a, b) => a.loadPriority - b.loadPriority
+    );
+
+    for (const Adapter of Adapters) {
+      const { CRUD, flavour: currentFlavour } = Adapter;
+      if (
+        Adapter.Events['app:access'] &&
+        !(await Adapter.Events['app:access']())
+      ) {
+        // skip the adapter if the user doesn't have access to it
+        const removed = metadata.filter(
+          meta => meta.flavour === currentFlavour
+        );
+        removed.forEach(meta => {
+          metadata.splice(metadata.indexOf(meta), 1);
+        });
+        continue;
+      }
+      try {
+        const item = await CRUD.list();
+        // remove the metadata that is not in the list
+        //  because we treat the workspace adapter as the source of truth
+        {
+          const removed = metadata.filter(
+            meta =>
+              meta.flavour === currentFlavour &&
+              !item.some(x => x.id === meta.id)
+          );
+          removed.forEach(meta => {
+            metadata.splice(metadata.indexOf(meta), 1);
+          });
+        }
+        // sort the metadata by the order of the list
+        if (metadata.length) {
+          item.sort((a, b) => {
+            return (
+              metadata.findIndex(x => x.id === a.id) -
+              metadata.findIndex(x => x.id === b.id)
+            );
+          });
+        }
+        metadata.push(
+          ...item.map(x => ({
+            id: x.id,
+            flavour: x.flavour,
+            version: WorkspaceVersion.DatabaseV3,
+          }))
+        );
+      } catch (e) {
+        console.error('list data error:', e);
+      }
+    }
+  }
+  const metadataMap = new Map(metadata.map(x => [x.id, x]));
+  // init workspace data
+  metadataMap.forEach((meta, id) => {
+    if (
+      meta.flavour === WorkspaceFlavour.AFFINE_CLOUD ||
+      meta.flavour === WorkspaceFlavour.LOCAL
+    ) {
+      getOrCreateWorkspace(id, meta.flavour);
+    } else {
+      throw new Error(`unknown flavour ${meta.flavour}`);
+    }
+  });
+  const result = Array.from(metadataMap.values());
+  console.info('metadata', result);
+  return result;
+};
+
 const rootWorkspacesMetadataPromiseAtom = atom<
   Promise<RootWorkspaceMetadata[]>
 >(async (get, { signal }) => {
-  const WorkspaceAdapters = get(workspaceAdaptersAtom);
-  assertExists(WorkspaceAdapters, 'workspace adapter should be defined');
   const primitiveMetadata = get(rootWorkspacesMetadataPrimitiveAtom);
   assertEquals(
     primitiveMetadata,
     null,
     'rootWorkspacesMetadataPrimitiveAtom should be null'
   );
-
-  if (environment.isServer) {
-    // return a promise in SSR to avoid the hydration mismatch
-    return Promise.resolve([]);
-  } else {
-    const metadata: RootWorkspaceMetadata[] = [];
-
-    // fixme(himself65): we might not need step 1
-    // step 1: try load metadata from localStorage
-    {
-      const loadFromLocalStorage = (): RootWorkspaceMetadata[] => {
-        // don't change this key,
-        // otherwise it will cause the data loss in the production
-        const primitiveMetadata = localStorage.getItem(METADATA_STORAGE_KEY);
-        if (primitiveMetadata) {
-          try {
-            const items = JSON.parse(primitiveMetadata) as z.infer<
-              typeof rootWorkspaceMetadataArraySchema
-            >;
-            rootWorkspaceMetadataArraySchema.parse(items);
-            return [...items];
-          } catch (e) {
-            console.error('cannot parse worksapce', e);
-          }
-          return [];
-        }
-        return [];
-      };
-
-      const maybeMetadata = loadFromLocalStorage();
-
-      // migration step, only data in `METADATA_STORAGE_KEY` will be migrated
-      if (
-        maybeMetadata.some(meta => !('version' in meta)) &&
-        !globalThis.$migrationDone
-      ) {
-        await new Promise<void>((resolve, reject) => {
-          signal.addEventListener('abort', () => reject(), { once: true });
-          window.addEventListener('migration-done', () => resolve(), {
-            once: true,
-          });
-        });
-      }
-
-      metadata.push(...loadFromLocalStorage());
-    }
-    // step 2: fetch from adapters
-    {
-      const Adapters = Object.values(WorkspaceAdapters).sort(
-        (a, b) => a.loadPriority - b.loadPriority
-      );
-
-      for (const Adapter of Adapters) {
-        const { CRUD, flavour: currentFlavour } = Adapter;
-        if (
-          Adapter.Events['app:access'] &&
-          !(await Adapter.Events['app:access']())
-        ) {
-          // skip the adapter if the user doesn't have access to it
-          continue;
-        }
-        try {
-          const item = await CRUD.list();
-          // remove the metadata that is not in the list
-          //  because we treat the workspace adapter as the source of truth
-          {
-            const removed = metadata.filter(
-              meta =>
-                meta.flavour === currentFlavour &&
-                !item.some(x => x.id === meta.id)
-            );
-            removed.forEach(meta => {
-              metadata.splice(metadata.indexOf(meta), 1);
-            });
-          }
-          // sort the metadata by the order of the list
-          if (metadata.length) {
-            item.sort((a, b) => {
-              return (
-                metadata.findIndex(x => x.id === a.id) -
-                metadata.findIndex(x => x.id === b.id)
-              );
-            });
-          }
-          metadata.push(
-            ...item.map(x => ({
-              id: x.id,
-              flavour: x.flavour,
-              version: WorkspaceVersion.DatabaseV3,
-            }))
-          );
-        } catch (e) {
-          console.error('list data error:', e);
-        }
-      }
-    }
-    const metadataMap = new Map(metadata.map(x => [x.id, x]));
-    // init workspace data
-    metadataMap.forEach((meta, id) => {
-      if (
-        meta.flavour === WorkspaceFlavour.AFFINE_CLOUD ||
-        meta.flavour === WorkspaceFlavour.LOCAL
-      ) {
-        getOrCreateWorkspace(id, meta.flavour);
-      } else {
-        throw new Error(`unknown flavour ${meta.flavour}`);
-      }
-    });
-    return Array.from(metadataMap.values());
-  }
+  return fetchMetadata(get, { signal });
 });
 
 type SetStateAction<Value> = Value | ((prev: Value) => Value);
@@ -246,6 +264,14 @@ export const rootWorkspacesMetadataAtom = atom<
   }
 );
 
+export const refreshRootMetadataAtom = atom(null, (get, set) => {
+  const abortController = new AbortController();
+  set(
+    rootWorkspacesMetadataPrimitiveAtom,
+    fetchMetadata(get, { signal: abortController.signal })
+  );
+});
+
 // blocksuite atoms,
 // each app should have only one block-hub in the same time
 export const rootBlockHubAtom = atom<Readonly<BlockHub> | null>(null);