fix: some login & enable affine cloud issues (#999)

Co-authored-by: himself65 <himself65@outlook.com>

apps/web/src/components/enable-workspace-modal/index.tsx

@@ -1,12 +1,10 @@
-import { Modal, ModalWrapper } from '@affine/component';
-import { IconButton } from '@affine/component';
-import { useTranslation } from '@affine/i18n';
 import { useAppState } from '@/providers/app-state-provider';
-import { useState } from 'react';
+import { IconButton, Modal, ModalWrapper, toast } from '@affine/component';
-import router from 'next/router';
+import { useTranslation } from '@affine/i18n';
-import { toast } from '@affine/component';
 import { CloseIcon } from '@blocksuite/icons';
-import { Header, Content, ContentTitle, StyleTips, StyleButton } from './style';
+import { useRouter } from 'next/router';
+import { useState } from 'react';
+import { Content, ContentTitle, Header, StyleButton, StyleTips } from './style';
 
 interface EnableWorkspaceModalProps {
   open: boolean;
@@ -20,6 +18,8 @@ export const EnableWorkspaceModal = ({
   const { t } = useTranslation();
   const { user, dataCenter, login, currentWorkspace } = useAppState();
   const [loading, setLoading] = useState(false);
+  const router = useRouter();
+
   return (
     <Modal open={open} onClose={onClose} data-testid="logout-modal">
       <ModalWrapper width={560} height={292}>
@@ -43,17 +43,19 @@ export const EnableWorkspaceModal = ({
               loading={loading}
               onClick={async () => {
                 setLoading(true);
-                if (!user) {
+                if (user || (await login())) {
-                  await login();
-                }
                   if (currentWorkspace) {
                     const workspace = await dataCenter.enableWorkspaceCloud(
                       currentWorkspace
                     );
-                  workspace &&
-                    router.push(`/workspace/${workspace.id}/setting`);
                     toast(t('Enabled success'));
+
+                    if (workspace) {
+                      router.push(`/workspace/${workspace.id}/setting`);
                     }
+                  }
+                }
+                setLoading(false);
               }}
             >
               {user ? t('Enable') : t('Sign in and Enable')}

apps/web/src/components/enable-workspace-modal/style.ts

@@ -6,6 +6,7 @@ export const Header = styled('div')({
   flexDirection: 'row-reverse',
   paddingRight: '10px',
   paddingTop: '10px',
+  flexShrink: 0,
 });
 
 export const Content = styled('div')({

apps/web/src/components/workspace-layout/index.tsx

@@ -4,10 +4,11 @@ import { useRouter } from 'next/router';
 import { StyledPage, StyledToolWrapper, StyledWrapper } from './styles';
 import { PropsWithChildren } from 'react';
 import useEnsureWorkspace from '@/hooks/use-ensure-workspace';
+import { PageLoading } from '../loading';
 
 export const WorkspaceDefender = ({ children }: PropsWithChildren) => {
   const { workspaceLoaded } = useEnsureWorkspace();
-  return <>{workspaceLoaded ? children : null}</>;
+  return <>{workspaceLoaded ? children : <PageLoading />}</>;
 };
 
 export const WorkspaceLayout = ({ children }: PropsWithChildren) => {

apps/web/src/components/workspace-modal/styles.ts

@@ -70,6 +70,7 @@ export const StyledCard = styled.div<{
 export const StyledFooter = styled('div')({
   height: '84px',
   padding: '0 40px',
+  flexShrink: 0,
   ...displayFlex('space-between', 'center'),
 });
 

apps/web/src/hooks/use-ensure-workspace.ts

@@ -1,11 +1,11 @@
-import { useState, useEffect } from 'react';
 import { useAppState } from '@/providers/app-state-provider';
 import { useRouter } from 'next/router';
+import { useEffect, useState } from 'react';
 // It is a fully effective hook
 // Cause it not just ensure workspace loaded, but also have router change.
 export const useEnsureWorkspace = () => {
   const [workspaceLoaded, setWorkspaceLoaded] = useState(false);
-  const { workspaceList, loadWorkspace, user } = useAppState();
+  const { dataCenter, loadWorkspace } = useAppState();
   const router = useRouter();
   const [activeWorkspaceId, setActiveWorkspaceId] = useState(
     router.query.workspaceId as string
@@ -14,15 +14,20 @@ export const useEnsureWorkspace = () => {
   // const defaultOutLineWorkspaceId = '99ce7eb7';
   // console.log(defaultOutLineWorkspaceId);
   useEffect(() => {
+    setWorkspaceLoaded(false);
+    let aborted = false;
+    const abortController = new AbortController();
+
+    const workspaceList = dataCenter.workspaces;
+    const workspaceId =
+      (router.query.workspaceId as string) || workspaceList[0]?.id;
+
     // If router.query.workspaceId is not in workspace list, jump to 404 page
     // If workspaceList is empty, we need to create a default workspace but not jump to 404
     if (
+      workspaceId &&
       workspaceList.length &&
-      // FIXME: router is not ready when this hook is called
+      workspaceList.findIndex(meta => meta.id.toString() === workspaceId) === -1
-      location.pathname.startsWith(`/workspace/${router.query.workspaceId}`) &&
-      workspaceList.findIndex(
-        meta => meta.id.toString() === router.query.workspaceId
-      ) === -1
     ) {
       router.push('/404');
       return;
@@ -36,13 +41,19 @@ export const useEnsureWorkspace = () => {
     //   router.push('/404');
     //   return;
     // }
-    const workspaceId =
+
-      (router.query.workspaceId as string) || workspaceList[0]?.id;
+    loadWorkspace.current(workspaceId, abortController.signal).then(unit => {
-    loadWorkspace.current(workspaceId).finally(() => {
+      if (!aborted && unit) {
         setWorkspaceLoaded(true);
-      setActiveWorkspaceId(activeWorkspaceId);
+        setActiveWorkspaceId(workspaceId);
+      }
     });
-  }, [loadWorkspace, router, user, workspaceList, activeWorkspaceId]);
+
+    return () => {
+      aborted = true;
+      abortController.abort();
+    };
+  }, [dataCenter, loadWorkspace, router]);
 
   return {
     workspaceLoaded,

apps/web/src/providers/app-state-provider/Provider.tsx

@@ -23,6 +23,7 @@ export const AppStateProvider = ({
   children,
 }: PropsWithChildren<AppStateContextProps>) => {
   const [appState, setAppState] = useState<AppStateValue>({} as AppStateValue);
+  const { dataCenter } = appState;
   const [blobState, setBlobState] = useState(false);
   const [userInfo, setUser] = useState<User | null>({} as User);
   useEffect(() => {
@@ -73,21 +74,20 @@ export const AppStateProvider = ({
   }, [appState]);
 
   useEffect(() => {
-    const { dataCenter } = appState;
     // FIXME: onWorkspacesChange should have dispose function
-    dataCenter?.onWorkspacesChange(
+    return dataCenter?.onWorkspacesChange(
       () => {
-        setAppState({
+        setAppState(_appState => ({
-          ...appState,
+          ..._appState,
           workspaceList: dataCenter.workspaces,
-        });
+        }));
       },
       { immediate: false }
     );
-  }, [appState]);
+  }, [dataCenter]);
 
   const loadPage = useRef<AppStateFunction['loadPage']>();
-  loadPage.current = (pageId: string) => {
+  loadPage.current = pageId => {
     const { currentWorkspace, currentPage } = appState;
     if (pageId === currentPage?.id) {
       return;
@@ -101,7 +101,7 @@ export const AppStateProvider = ({
 
   const loadWorkspace: AppStateFunction['loadWorkspace'] =
     useRef() as AppStateFunction['loadWorkspace'];
-  loadWorkspace.current = async (workspaceId: string) => {
+  loadWorkspace.current = async (workspaceId, abort) => {
     const { dataCenter, workspaceList, currentWorkspace } = appState;
     if (!workspaceList.find(v => v.id.toString() === workspaceId)) {
       return null;
@@ -110,7 +110,21 @@ export const AppStateProvider = ({
       return currentWorkspace;
     }
 
+    let aborted = false;
+
+    const onAbort = () => {
+      aborted = true;
+    };
+
+    abort?.addEventListener('abort', onAbort);
+
     const workspace = (await dataCenter.loadWorkspace(workspaceId)) ?? null;
+
+    if (aborted) {
+      // do not update state if aborted
+      return null;
+    }
+
     let isOwner;
     if (workspace?.provider === 'local') {
       // isOwner is useful only in the cloud
@@ -132,6 +146,8 @@ export const AppStateProvider = ({
       isOwner,
     });
 
+    abort?.removeEventListener('abort', onAbort);
+
     return workspace;
   };
 
@@ -174,15 +190,19 @@ export const AppStateProvider = ({
 
   const login = async () => {
     const { dataCenter } = appState;
+    try {
       await dataCenter.login();
-
       const user = (await dataCenter.getUserInfo()) as User;
       if (!user) {
         throw new Error('User info not found');
       }
       setUser(user);
       return user;
+    } catch (error) {
+      return null; // login failed
+    }
   };
+
   const logout = async () => {
     const { dataCenter } = appState;
     await dataCenter.logout();

apps/web/src/providers/app-state-provider/interface.ts

@@ -34,11 +34,11 @@ export type AppStateFunction = {
   setBlockHub: MutableRefObject<(BlockHub: BlockHub) => void>;
 
   loadWorkspace: MutableRefObject<
-    (workspaceId: string) => Promise<WorkspaceUnit | null>
+    (workspaceId: string, abort?: AbortSignal) => Promise<WorkspaceUnit | null>
   >;
   loadPage: (pageId: string) => void;
 
-  login: () => Promise<User>;
+  login: () => Promise<User | null>;
   logout: () => Promise<void>;
 };
 

apps/web/src/store/globalModal/index.tsx

@@ -1,5 +1,11 @@
 import type React from 'react';
-import { createContext, useCallback, useContext, useMemo } from 'react';
+import {
+  createContext,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+} from 'react';
 import { createStore, useStore } from 'zustand';
 import { combine, subscribeWithSelector } from 'zustand/middleware';
 import { UseBoundStore } from 'zustand/react';
@@ -9,6 +15,7 @@ import QuickSearch from '@/components/quick-search';
 import { LoginModal } from '@/components/login-modal';
 import ImportModal from '@/components/import';
 import { EnableWorkspaceModal } from '@/components/enable-workspace-modal';
+import { useRouter } from 'next/router';
 
 export type ModalState = {
   contact: boolean;
@@ -28,19 +35,19 @@ export type ModalActions = {
   triggerEnableWorkspaceModal: () => void;
 };
 
-const create = () =>
+const defaultModalState: ModalState = {
-  createStore(
-    subscribeWithSelector(
-      combine<ModalState, ModalActions>(
-        {
   contact: false,
   shortcuts: false,
   quickSearch: false,
   import: false,
   login: false,
   enableWorkspace: false,
-        },
+};
-        set => ({
+
+const create = () =>
+  createStore(
+    subscribeWithSelector(
+      combine<ModalState, ModalActions>({ ...defaultModalState }, set => ({
         triggerShortcutsModal: () => {
           set(({ shortcuts }) => ({
             shortcuts: !shortcuts,
@@ -71,8 +78,7 @@ const create = () =>
             enableWorkspace: !enableWorkspace,
           }));
         },
-        })
+      }))
-      )
     )
   );
 type Store = ReturnType<typeof create>;
@@ -81,9 +87,20 @@ const ModalContext = createContext<Store | null>(null);
 
 export const useModalApi = () => {
   const api = useContext(ModalContext);
+
   if (!api) {
     throw new Error('cannot find modal context');
   }
+
+  const router = useRouter();
+
+  useEffect(() => {
+    router.events.on('routeChangeStart', () => {
+      // normally modal should be closed when route change
+      api.setState(defaultModalState);
+    });
+  }, [api, router.events]);
+
   return api;
 };
 

packages/component/src/ui/modal/ModalWrapper.tsx

@@ -13,6 +13,7 @@ export const ModalWrapper = styled.div<{
     backgroundColor: theme.colors.popoverBackground,
     borderRadius: '24px',
     position: 'relative',
+    maxHeight: 'calc(100vh - 32px)',
   };
 });
 

packages/data-center/src/datacenter.ts

@@ -1,9 +1,6 @@
 import { WorkspaceUnitCollection } from './workspace-unit-collection.js';
 import type { WorkspaceUnitCollectionChangeEvent } from './workspace-unit-collection';
-import {
+import { Workspace as BlocksuiteWorkspace } from '@blocksuite/store';
-  StoreOptions,
-  Workspace as BlocksuiteWorkspace,
-} from '@blocksuite/store';
 import type {
   BaseProvider,
   CreateWorkspaceInfoParams,
@@ -127,12 +124,12 @@ export class DataCenter {
    * get a new workspace only has room id
    * @param {string} workspaceId workspace id
    */
-  private _getBlocksuiteWorkspace(workspaceId: string, params: StoreOptions) {
+  private _getBlocksuiteWorkspace(workspaceId: string) {
     // const workspaceInfo = this._workspaceUnitCollection.find(workspaceId);
     // assert(workspaceInfo, 'Workspace not found');
     return (
       // this._workspaceInstances.get(workspaceId) ||
-      createBlocksuiteWorkspace(workspaceId, params)
+      createBlocksuiteWorkspace(workspaceId)
     );
   }
 
@@ -174,14 +171,7 @@ export class DataCenter {
     assert(provider, `provide '${workspaceUnit.provider}' is not registered`);
     this._logger(`Loading ${workspaceUnit.provider} workspace: `, workspaceId);
 
-    const params: StoreOptions = {};
+    const workspace = this._getBlocksuiteWorkspace(workspaceId);
-    if (provider.id === 'affine') {
-      params.blobOptionsGetter = (k: string) =>
-        ({ api: '/api/workspace', token: provider.getToken() }[k]);
-    } else {
-      params.blobOptionsGetter = (k: string) => undefined;
-    }
-    const workspace = this._getBlocksuiteWorkspace(workspaceId, params);
     this._workspaceInstances.set(workspaceId, workspace);
     await provider.warpWorkspace(workspace);
     this._workspaceUnitCollection.workspaces.forEach(workspaceUnit => {
@@ -197,7 +187,7 @@ export class DataCenter {
     // FIXME: hard code for public workspace
     const provider = this.providerMap.get('affine');
     assert(provider);
-    const blocksuiteWorkspace = this._getBlocksuiteWorkspace(workspaceId, {});
+    const blocksuiteWorkspace = this._getBlocksuiteWorkspace(workspaceId);
     await provider.loadPublicWorkspace(blocksuiteWorkspace);
 
     const workspaceUnitForPublic = new WorkspaceUnit({
@@ -231,7 +221,7 @@ export class DataCenter {
    * listen workspaces list change
    * @param {Function} callback callback function
    */
-  public async onWorkspacesChange(
+  public onWorkspacesChange(
     callback: (workspaces: WorkspaceUnitCollectionChangeEvent) => void,
     { immediate = true }: { immediate?: boolean } = {}
   ) {
@@ -241,6 +231,9 @@ export class DataCenter {
       });
     }
     this._workspaceUnitCollection.on('change', callback);
+    return () => {
+      this._workspaceUnitCollection.off('change', callback);
+    };
   }
 
   /**

packages/data-center/src/provider/affine/affine.ts

@@ -15,9 +15,10 @@ import {
   loadWorkspaceUnit,
   createWorkspaceUnit,
   migrateBlobDB,
+  createBlocksuiteWorkspaceWithAuth,
 } from './utils.js';
 import { WorkspaceUnit } from '../../workspace-unit.js';
-import { createBlocksuiteWorkspace, applyUpdate } from '../../utils/index.js';
+import { applyUpdate } from '../../utils/index.js';
 import type { SyncMode } from '../../workspace-unit';
 import { MessageCenter } from '../../message/index.js';
 
@@ -43,7 +44,10 @@ export class AffineProvider extends BaseProvider {
   private _channel?: WebsocketClient;
   private _refreshToken?: string;
   // private _idbMap: Map<string, IndexedDBProvider> = new Map();
-  private _workspaceLoadingQueue: Set<string> = new Set();
+  private _workspaceLoadingQueue: Map<string, Promise<WorkspaceUnit>> =
+    new Map();
+
+  private _workspaces$: Promise<Workspace[]> | undefined;
 
   constructor({ apis, ...params }: AffineProviderConstructorParams) {
     super(params);
@@ -141,11 +145,10 @@ export class AffineProvider extends BaseProvider {
           this._workspaces.update(id, workspace);
         } else {
           if (!this._workspaceLoadingQueue.has(id)) {
-            const workspaceUnit = await loadWorkspaceUnit(
+            const p = loadWorkspaceUnit({ id, ...workspace }, this._apis);
-              { id, ...workspace },
+            this._workspaceLoadingQueue.set(id, p);
-              this._apis
+            newlyCreatedWorkspaces.push(await p);
-            );
+            this._workspaceLoadingQueue.delete(id);
-            newlyCreatedWorkspaces.push(workspaceUnit);
           }
         }
       } else {
@@ -201,6 +204,9 @@ export class AffineProvider extends BaseProvider {
   }
 
   override async warpWorkspace(workspace: BlocksuiteWorkspace) {
+    workspace.setGettingBlobOptions(
+      (k: string) => ({ api: '/api/workspace', token: this.getToken() }[k])
+    );
     // FIXME: if add indexedDB cache in the future, can remove following line.
     await this._applyCloudUpdates(workspace);
     const { room } = workspace;
@@ -229,11 +235,22 @@ export class AffineProvider extends BaseProvider {
     if (!this._apis.auth.isLogin) {
       return [];
     }
-    const workspacesList = await this._apis.getWorkspaces();
+
+    // cache workspaces and workspaceUnits results so that simultaneous calls
+    // to loadWorkspaces will not cause multiple requests
+    if (!this._workspaces$) {
+      this._workspaces$ = this._apis.getWorkspaces();
+    }
+
+    const workspacesList = await this._workspaces$;
     const workspaceUnits = await Promise.all(
-      workspacesList.map(w => {
+      workspacesList.map(async w => {
-        this._workspaceLoadingQueue.add(w.id);
+        let p = this._workspaceLoadingQueue.get(w.id);
-        return loadWorkspaceUnit(
+        if (!p) {
+          // may only need to load the primary one instead of all of them?
+          // it will take a long time to load all of the workspaces
+          // at least we shall use p-map to load them in chunks
+          p = loadWorkspaceUnit(
             {
               id: w.id,
               name: '',
@@ -245,11 +262,19 @@ export class AffineProvider extends BaseProvider {
               syncMode: 'core',
             },
             this._apis
-        ).finally(() => {
+          );
+          this._workspaceLoadingQueue.set(w.id, p);
+        }
+        const workspaceUnit = await p;
         this._workspaceLoadingQueue.delete(w.id);
-        });
+
+        return workspaceUnit;
       })
     );
+
+    // release cache
+    this._workspaces$ = undefined;
+
     this._workspaces.add(workspaceUnits);
     return workspaceUnits;
   }
@@ -420,7 +445,7 @@ export class AffineProvider extends BaseProvider {
     });
     await migrateBlobDB(workspaceUnit.id, id);
 
-    const blocksuiteWorkspace = createBlocksuiteWorkspace(id);
+    const blocksuiteWorkspace = await createBlocksuiteWorkspaceWithAuth(id);
     assert(workspaceUnit.blocksuiteWorkspace);
     await applyUpdate(
       blocksuiteWorkspace,

packages/data-center/src/provider/affine/apis/auth.ts

@@ -181,6 +181,7 @@ export const auth = new Auth();
 
 export const getAuthorizer = () => {
   let _firebaseAuth: FirebaseAuth | null = null;
+  const logger = getLogger('authorizer');
 
   // getAuth will send requests on calling thus we can lazy init it
   const getAuth = () => {
@@ -200,8 +201,7 @@ export const getAuthorizer = () => {
       }
       return _firebaseAuth;
     } catch (error) {
-      getLogger('getAuthorizer')(error);
+      logger(error);
-      console.error('getAuthorizer', error);
       return null;
     }
   };

packages/data-center/src/provider/affine/apis/request.ts

@@ -11,6 +11,8 @@ const _sendMessage = messageCenter.getMessageSender('affine');
 export const bareClient: KyInstance = ky.extend({
   prefixUrl: '/',
   retry: 1,
+  // todo: report timeout error
+  timeout: 60000,
   hooks: {
     // afterResponse: [
     //   async (_request, _options, response) => {

packages/data-center/src/provider/affine/utils.ts

@@ -1,17 +1,31 @@
 import { WorkspaceUnit } from '../../workspace-unit.js';
 import type { WorkspaceUnitCtorParams } from '../../workspace-unit';
-import { createBlocksuiteWorkspace } from '../../utils/index.js';
+import { createBlocksuiteWorkspace as _createBlocksuiteWorkspace } from '../../utils/index.js';
 import type { Apis } from './apis';
 import { setDefaultAvatar } from '../utils.js';
 import { applyUpdate } from '../../utils/index.js';
 import { getDatabase } from './idb-kv.js';
+import { auth } from './apis/auth.js';
+
+export const createBlocksuiteWorkspaceWithAuth = async (id: string) => {
+  if (auth.isExpired && auth.isLogin) {
+    await auth.refreshToken();
+  }
+  return _createBlocksuiteWorkspace(id, {
+    blobOptionsGetter: (k: string) =>
+      // token could be expired
+      ({ api: '/api/workspace', token: auth.token }[k]),
+  });
+};
 
 export const loadWorkspaceUnit = async (
   params: WorkspaceUnitCtorParams,
   apis: Apis
 ) => {
   const workspaceUnit = new WorkspaceUnit(params);
-  const blocksuiteWorkspace = createBlocksuiteWorkspace(workspaceUnit.id);
+  const blocksuiteWorkspace = await createBlocksuiteWorkspaceWithAuth(
+    workspaceUnit.id
+  );
 
   const updates = await apis.downloadWorkspace(
     workspaceUnit.id,
@@ -43,7 +57,9 @@ export const loadWorkspaceUnit = async (
 export const createWorkspaceUnit = async (params: WorkspaceUnitCtorParams) => {
   const workspaceUnit = new WorkspaceUnit(params);
 
-  const blocksuiteWorkspace = createBlocksuiteWorkspace(workspaceUnit.id);
+  const blocksuiteWorkspace = await createBlocksuiteWorkspaceWithAuth(
+    workspaceUnit.id
+  );
 
   blocksuiteWorkspace.meta.setName(workspaceUnit.name);
   if (!workspaceUnit.avatar) {

packages/data-center/src/workspace-unit-collection.ts

@@ -37,6 +37,13 @@ export class WorkspaceUnitCollection {
     this._events.on(type, callback);
   }
 
+  public off(
+    type: 'change',
+    callback: (event: WorkspaceUnitCollectionChangeEvent) => void
+  ) {
+    this._events.off(type, callback);
+  }
+
   public once(
     type: 'change',
     callback: (event: WorkspaceUnitCollectionChangeEvent) => void

packages/data-center/src/workspace-unit.ts

@@ -1,4 +1,7 @@
-import { Workspace as BlocksuiteWorkspace } from '@blocksuite/store';
+import {
+  BlobOptionsGetter,
+  Workspace as BlocksuiteWorkspace,
+} from '@blocksuite/store';
 import type { User } from './types';
 
 export type SyncMode = 'all' | 'core';
@@ -13,6 +16,7 @@ export interface WorkspaceUnitCtorParams {
   provider: string;
   syncMode: SyncMode;
 
+  blobOptionsGetter?: BlobOptionsGetter;
   blocksuiteWorkspace?: BlocksuiteWorkspace | null;
 }
 
@@ -60,9 +64,18 @@ export class WorkspaceUnit {
 
   update(params: UpdateWorkspaceUnitParams) {
     Object.assign(this, params);
+    if (params.blocksuiteWorkspace) {
+      this.setBlocksuiteWorkspace(params.blocksuiteWorkspace);
+    }
+    if (params.blobOptionsGetter && this.blocksuiteWorkspace) {
+      this.blocksuiteWorkspace.setGettingBlobOptions(params.blobOptionsGetter);
+    }
   }
 
-  toJSON(): Omit<WorkspaceUnitCtorParams, 'blocksuiteWorkspace'> {
+  toJSON(): Omit<
+    WorkspaceUnitCtorParams,
+    'blocksuiteWorkspace' | 'blobOptionsGetter'
+  > {
     return {
       id: this.id,
       name: this.name,