feat(core): desktop multiple server support (#8979)

packages/frontend/core/src/modules/cloud/entities/server.ts

@@ -11,6 +11,7 @@ import {
 import { EMPTY, exhaustMap, map, mergeMap } from 'rxjs';
 
 import { ServerScope } from '../scopes/server';
+import { AuthService } from '../services/auth';
 import { FetchService } from '../services/fetch';
 import { GraphQLService } from '../services/graphql';
 import { ServerConfigStore } from '../stores/server-config';
@@ -34,6 +35,9 @@ export class Server extends Entity<{
   readonly serverConfigStore = this.scope.framework.get(ServerConfigStore);
   readonly fetch = this.scope.framework.get(FetchService).fetch;
   readonly gql = this.scope.framework.get(GraphQLService).gql;
+  get account$() {
+    return this.scope.framework.get(AuthService).session.account$;
+  }
   readonly serverMetadata = this.props.serverMetadata;
 
   constructor(private readonly serverListStore: ServerListStore) {
@@ -68,7 +72,7 @@ export class Server extends Entity<{
   readonly revalidateConfig = effect(
     exhaustMap(() => {
       return fromPromise(signal =>
-        this.serverConfigStore.fetchServerConfig(signal)
+        this.serverConfigStore.fetchServerConfig(this.baseUrl, signal)
       ).pipe(
         backoffRetry({
           count: Infinity,

packages/frontend/core/src/modules/cloud/entities/session.ts

@@ -68,7 +68,7 @@ export class AuthSession extends Entity {
 
   revalidate = effect(
     exhaustMapWithTrailing(() =>
-      fromPromise(this.getSession()).pipe(
+      fromPromise(() => this.getSession()).pipe(
         backoffRetry({
           count: Infinity,
         }),

packages/frontend/core/src/modules/cloud/events/account-changed.ts

@@ -0,0 +1,7 @@
+import { createEvent } from '@toeverything/infra';
+
+import type { AuthAccountInfo } from '../entities/session';
+
+export const AccountChanged = createEvent<AuthAccountInfo | null>(
+  'AccountChanged'
+);

packages/frontend/core/src/modules/cloud/events/account-logged-in.ts

@@ -0,0 +1,5 @@
+import { createEvent } from '@toeverything/infra';
+
+import type { AuthAccountInfo } from '../entities/session';
+
+export const AccountLoggedIn = createEvent<AuthAccountInfo>('AccountLoggedIn');

packages/frontend/core/src/modules/cloud/events/account-logged-out.ts

@@ -0,0 +1,6 @@
+import { createEvent } from '@toeverything/infra';
+
+import type { AuthAccountInfo } from '../entities/session';
+
+export const AccountLoggedOut =
+  createEvent<AuthAccountInfo>('AccountLoggedOut');

packages/frontend/core/src/modules/cloud/events/server-initialized.ts

@@ -0,0 +1,5 @@
+import { createEvent } from '@toeverything/infra';
+
+import type { Server } from '../entities/server';
+
+export const ServerInitialized = createEvent<Server>('ServerInitialized');

packages/frontend/core/src/modules/cloud/events/server-started.ts

@@ -0,0 +1,3 @@
+import { createEvent } from '@toeverything/infra';
+
+export const ServerStarted = createEvent('ServerStarted');

packages/frontend/core/src/modules/cloud/index.ts

@@ -7,10 +7,14 @@ export {
   isNetworkError,
   NetworkError,
 } from './error';
+export { AccountChanged } from './events/account-changed';
+export { AccountLoggedIn } from './events/account-logged-in';
+export { AccountLoggedOut } from './events/account-logged-out';
+export { ServerInitialized } from './events/server-initialized';
 export { RawFetchProvider } from './provider/fetch';
 export { ValidatorProvider } from './provider/validator';
 export { WebSocketAuthProvider } from './provider/websocket-auth';
-export { AccountChanged, AuthService } from './services/auth';
+export { AuthService } from './services/auth';
 export { CaptchaService } from './services/captcha';
 export { DefaultServerService } from './services/default-server';
 export { EventSourceService } from './services/eventsource';
@@ -25,6 +29,7 @@ export { UserFeatureService } from './services/user-feature';
 export { UserQuotaService } from './services/user-quota';
 export { WebSocketService } from './services/websocket';
 export { WorkspaceServerService } from './services/workspace-server';
+export type { ServerConfig } from './types';
 
 import {
   DocScope,
@@ -79,9 +84,10 @@ import { UserQuotaStore } from './stores/user-quota';
 export function configureCloudModule(framework: Framework) {
   framework
     .impl(RawFetchProvider, DefaultRawFetchProvider)
-    .service(ServersService, [ServerListStore])
+    .service(ServersService, [ServerListStore, ServerConfigStore])
     .service(DefaultServerService, [ServersService])
     .store(ServerListStore, [GlobalStateService])
+    .store(ServerConfigStore, [RawFetchProvider])
     .entity(Server, [ServerListStore])
     .scope(ServerScope)
     .service(ServerService, [ServerScope])
@@ -97,7 +103,6 @@ export function configureCloudModule(framework: Framework) {
           f.getOptional(WebSocketAuthProvider)
         )
     )
-    .store(ServerConfigStore, [GraphQLService])
     .service(CaptchaService, f => {
       return new CaptchaService(
         f.get(ServerService),
@@ -106,7 +111,12 @@ export function configureCloudModule(framework: Framework) {
       );
     })
     .service(AuthService, [FetchService, AuthStore, UrlService])
-    .store(AuthStore, [FetchService, GraphQLService, GlobalState])
+    .store(AuthStore, [
+      FetchService,
+      GraphQLService,
+      GlobalState,
+      ServerService,
+    ])
     .entity(AuthSession, [AuthStore])
     .service(SubscriptionService, [SubscriptionStore])
     .store(SubscriptionStore, [
@@ -132,12 +142,13 @@ export function configureCloudModule(framework: Framework) {
     .store(UserFeatureStore, [GraphQLService])
     .service(InvoicesService)
     .store(InvoicesStore, [GraphQLService])
-    .entity(Invoices, [InvoicesStore])
+    .entity(Invoices, [InvoicesStore]);
+
+  framework
     .scope(WorkspaceScope)
+    .service(WorkspaceServerService)
     .scope(DocScope)
     .service(CloudDocMetaService)
     .entity(CloudDocMeta, [CloudDocMetaStore, DocService, GlobalCache])
-    .store(CloudDocMetaStore, [GraphQLService]);
-
-  framework.scope(WorkspaceScope).service(WorkspaceServerService);
+    .store(CloudDocMetaStore, [WorkspaceServerService]);
 }

packages/frontend/core/src/modules/cloud/services/auth.ts

@@ -1,18 +1,16 @@
 import { AIProvider } from '@affine/core/blocksuite/presets/ai';
 import type { OAuthProviderType } from '@affine/graphql';
 import { track } from '@affine/track';
-import {
-  ApplicationFocused,
-  ApplicationStarted,
-  createEvent,
-  OnEvent,
-  Service,
-} from '@toeverything/infra';
+import { ApplicationFocused, OnEvent, Service } from '@toeverything/infra';
 import { distinctUntilChanged, map, skip } from 'rxjs';
 
 import type { UrlService } from '../../url';
 import { type AuthAccountInfo, AuthSession } from '../entities/session';
 import { BackendError } from '../error';
+import { AccountChanged } from '../events/account-changed';
+import { AccountLoggedIn } from '../events/account-logged-in';
+import { AccountLoggedOut } from '../events/account-logged-out';
+import { ServerStarted } from '../events/server-started';
 import type { AuthStore } from '../stores/auth';
 import type { FetchService } from './fetch';
 
@@ -26,18 +24,8 @@ function toAIUserInfo(account: AuthAccountInfo | null) {
   };
 }
 
-// Emit when account changed
-export const AccountChanged = createEvent<AuthAccountInfo | null>(
-  'AccountChanged'
-);
-
-export const AccountLoggedIn = createEvent<AuthAccountInfo>('AccountLoggedIn');
-
-export const AccountLoggedOut =
-  createEvent<AuthAccountInfo>('AccountLoggedOut');
-
-@OnEvent(ApplicationStarted, e => e.onApplicationStart)
 @OnEvent(ApplicationFocused, e => e.onApplicationFocused)
+@OnEvent(ServerStarted, e => e.onServerStarted)
 export class AuthService extends Service {
   session = this.framework.createEntity(AuthSession);
 
@@ -74,7 +62,7 @@ export class AuthService extends Service {
       });
   }
 
-  private onApplicationStart() {
+  private onServerStarted() {
     this.session.revalidate();
   }
 

packages/frontend/core/src/modules/cloud/services/captcha.ts

@@ -7,7 +7,7 @@ import {
   onStart,
   Service,
 } from '@toeverything/infra';
-import { EMPTY, exhaustMap, mergeMap } from 'rxjs';
+import { EMPTY, exhaustMap, mergeMap, switchMap } from 'rxjs';
 
 import type { ValidatorProvider } from '../provider/validator';
 import type { FetchService } from './fetch';
@@ -61,10 +61,12 @@ export class CaptchaService extends Service {
         mergeMap(({ challenge, token }) => {
           this.verifyToken$.next(token);
           this.challenge$.next(challenge);
+          this.resetAfter5min();
           return EMPTY;
         }),
         catchErrorInto(this.error$),
         onStart(() => {
+          this.challenge$.next(undefined);
           this.verifyToken$.next(undefined);
           this.isLoading$.next(true);
         }),
@@ -72,4 +74,22 @@ export class CaptchaService extends Service {
       );
     })
   );
+
+  resetAfter5min = effect(
+    switchMap(() => {
+      return fromPromise(async () => {
+        await new Promise(resolve => {
+          setTimeout(resolve, 1000 * 60 * 5);
+        });
+        return true;
+      }).pipe(
+        mergeMap(_ => {
+          this.challenge$.next(undefined);
+          this.verifyToken$.next(undefined);
+          this.isLoading$.next(false);
+          return EMPTY;
+        })
+      );
+    })
+  );
 }

packages/frontend/core/src/modules/cloud/services/servers.ts

@@ -1,12 +1,20 @@
+import { Unreachable } from '@affine/env/constant';
 import { LiveData, ObjectPool, Service } from '@toeverything/infra';
-import { finalize, of, switchMap } from 'rxjs';
+import { nanoid } from 'nanoid';
+import { Observable, switchMap } from 'rxjs';
 
 import { Server } from '../entities/server';
+import { ServerInitialized } from '../events/server-initialized';
+import { ServerStarted } from '../events/server-started';
+import type { ServerConfigStore } from '../stores/server-config';
 import type { ServerListStore } from '../stores/server-list';
 import type { ServerConfig, ServerMetadata } from '../types';
 
 export class ServersService extends Service {
-  constructor(private readonly serverListStore: ServerListStore) {
+  constructor(
+    private readonly serverListStore: ServerListStore,
+    private readonly serverConfigStore: ServerConfigStore
+  ) {
     super();
   }
 
@@ -21,17 +29,21 @@ export class ServersService extends Service {
           const server = this.framework.createEntity(Server, {
             serverMetadata: metadata,
           });
+          server.revalidateConfig();
+          this.eventBus.emit(ServerInitialized, server);
+          server.scope.eventBus.emit(ServerStarted, server);
           const ref = this.serverPool.put(metadata.id, server);
           return ref;
         });
 
-        return of(refs.map(ref => ref.obj)).pipe(
-          finalize(() => {
+        return new Observable<Server[]>(subscribe => {
+          subscribe.next(refs.map(ref => ref.obj));
+          return () => {
             refs.forEach(ref => {
               ref.release();
             });
-          })
-        );
+          };
+        });
       })
     ),
     [] as any
@@ -52,4 +64,43 @@ export class ServersService extends Service {
   addServer(metadata: ServerMetadata, config: ServerConfig) {
     this.serverListStore.addServer(metadata, config);
   }
+
+  removeServer(id: string) {
+    this.serverListStore.removeServer(id);
+  }
+
+  async addServerByBaseUrl(baseUrl: string) {
+    const config = await this.serverConfigStore.fetchServerConfig(baseUrl);
+    const id = nanoid();
+    this.serverListStore.addServer(
+      { id, baseUrl },
+      {
+        credentialsRequirement: config.credentialsRequirement,
+        features: config.features,
+        oauthProviders: config.oauthProviders,
+        serverName: config.name,
+        type: config.type,
+        initialized: config.initialized,
+        version: config.version,
+      }
+    );
+  }
+
+  getServerByBaseUrl(baseUrl: string) {
+    return this.servers$.value.find(s => s.baseUrl === baseUrl);
+  }
+
+  async addOrGetServerByBaseUrl(baseUrl: string) {
+    const server = this.getServerByBaseUrl(baseUrl);
+    if (server) {
+      return server;
+    } else {
+      await this.addServerByBaseUrl(baseUrl);
+      const server = this.getServerByBaseUrl(baseUrl);
+      if (!server) {
+        throw new Unreachable();
+      }
+      return server;
+    }
+  }
 }

packages/frontend/core/src/modules/cloud/services/subscription.ts

@@ -4,8 +4,8 @@ import { OnEvent, Service } from '@toeverything/infra';
 
 import { Subscription } from '../entities/subscription';
 import { SubscriptionPrices } from '../entities/subscription-prices';
+import { AccountChanged } from '../events/account-changed';
 import type { SubscriptionStore } from '../stores/subscription';
-import { AccountChanged } from './auth';
 
 @OnEvent(AccountChanged, e => e.onAccountChanged)
 export class SubscriptionService extends Service {

packages/frontend/core/src/modules/cloud/services/user-copilot-quota.ts

@@ -1,7 +1,7 @@
 import { OnEvent, Service } from '@toeverything/infra';
 
 import { UserCopilotQuota } from '../entities/user-copilot-quota';
-import { AccountChanged } from './auth';
+import { AccountChanged } from '../events/account-changed';
 
 @OnEvent(AccountChanged, e => e.onAccountChanged)
 export class UserCopilotQuotaService extends Service {

packages/frontend/core/src/modules/cloud/services/user-feature.ts

@@ -1,7 +1,7 @@
 import { OnEvent, Service } from '@toeverything/infra';
 
 import { UserFeature } from '../entities/user-feature';
-import { AccountChanged } from './auth';
+import { AccountChanged } from '../events/account-changed';
 
 @OnEvent(AccountChanged, e => e.onAccountChanged)
 export class UserFeatureService extends Service {

packages/frontend/core/src/modules/cloud/services/user-quota.ts

@@ -2,7 +2,7 @@ import { mixpanel } from '@affine/track';
 import { OnEvent, Service } from '@toeverything/infra';
 
 import { UserQuota } from '../entities/user-quota';
-import { AccountChanged } from './auth';
+import { AccountChanged } from '../events/account-changed';
 
 @OnEvent(AccountChanged, e => e.onAccountChanged)
 export class UserQuotaService extends Service {

packages/frontend/core/src/modules/cloud/services/websocket.ts

@@ -1,9 +1,9 @@
 import { ApplicationStarted, OnEvent, Service } from '@toeverything/infra';
 import { Manager } from 'socket.io-client';
 
+import { AccountChanged } from '../events/account-changed';
 import type { WebSocketAuthProvider } from '../provider/websocket-auth';
 import type { AuthService } from './auth';
-import { AccountChanged } from './auth';
 import type { ServerService } from './server';
 
 @OnEvent(AccountChanged, e => e.update)

packages/frontend/core/src/modules/cloud/stores/auth.ts

@@ -9,6 +9,7 @@ import { Store } from '@toeverything/infra';
 import type { AuthSessionInfo } from '../entities/session';
 import type { FetchService } from '../services/fetch';
 import type { GraphQLService } from '../services/graphql';
+import type { ServerService } from '../services/server';
 
 export interface AccountProfile {
   id: string;
@@ -23,21 +24,26 @@ export class AuthStore extends Store {
   constructor(
     private readonly fetchService: FetchService,
     private readonly gqlService: GraphQLService,
-    private readonly globalState: GlobalState
+    private readonly globalState: GlobalState,
+    private readonly serverService: ServerService
   ) {
     super();
   }
 
   watchCachedAuthSession() {
-    return this.globalState.watch<AuthSessionInfo>('affine-cloud-auth');
+    return this.globalState.watch<AuthSessionInfo>(
+      `${this.serverService.server.id}-auth`
+    );
   }
 
   getCachedAuthSession() {
-    return this.globalState.get<AuthSessionInfo>('affine-cloud-auth');
+    return this.globalState.get<AuthSessionInfo>(
+      `${this.serverService.server.id}-auth`
+    );
   }
 
   setCachedAuthSession(session: AuthSessionInfo | null) {
-    this.globalState.set('affine-cloud-auth', session);
+    this.globalState.set(`${this.serverService.server.id}-auth`, session);
   }
 
   async fetchSession() {
@@ -99,6 +105,7 @@ export class AuthStore extends Store {
     const data = (await res.json()) as {
       registered: boolean;
       hasPassword: boolean;
+      magicLink: boolean;
     };
 
     return data;

packages/frontend/core/src/modules/cloud/stores/cloud-doc-meta.ts

@@ -2,10 +2,10 @@ import { getWorkspacePageMetaByIdQuery } from '@affine/graphql';
 import { Store } from '@toeverything/infra';
 
 import { type CloudDocMetaType } from '../entities/cloud-doc-meta';
-import type { GraphQLService } from '../services/graphql';
+import type { WorkspaceServerService } from '../services/workspace-server';
 
 export class CloudDocMetaStore extends Store {
-  constructor(private readonly gqlService: GraphQLService) {
+  constructor(private readonly workspaceServerService: WorkspaceServerService) {
     super();
   }
 
@@ -14,7 +14,10 @@ export class CloudDocMetaStore extends Store {
     docId: string,
     abortSignal?: AbortSignal
   ): Promise<CloudDocMetaType> {
-    const serverConfigData = await this.gqlService.gql({
+    if (!this.workspaceServerService.server) {
+      throw new Error('Server not found');
+    }
+    const serverConfigData = await this.workspaceServerService.server.gql({
       query: getWorkspacePageMetaByIdQuery,
       variables: { id: workspaceId, pageId: docId },
       context: {

packages/frontend/core/src/modules/cloud/stores/server-config.ts

@@ -1,4 +1,5 @@
 import {
+  gqlFetcherFactory,
   type OauthProvidersQuery,
   oauthProvidersQuery,
   type ServerConfigQuery,
@@ -7,27 +8,32 @@ import {
 } from '@affine/graphql';
 import { Store } from '@toeverything/infra';
 
-import type { GraphQLService } from '../services/graphql';
+import type { RawFetchProvider } from '../provider/fetch';
 
 export type ServerConfigType = ServerConfigQuery['serverConfig'] &
   OauthProvidersQuery['serverConfig'];
 
 export class ServerConfigStore extends Store {
-  constructor(private readonly gqlService: GraphQLService) {
+  constructor(private readonly fetcher: RawFetchProvider) {
     super();
   }
 
   async fetchServerConfig(
+    serverBaseUrl: string,
     abortSignal?: AbortSignal
   ): Promise<ServerConfigType> {
-    const serverConfigData = await this.gqlService.gql({
+    const gql = gqlFetcherFactory(
+      `${serverBaseUrl}/graphql`,
+      this.fetcher.fetch
+    );
+    const serverConfigData = await gql({
       query: serverConfigQuery,
       context: {
         signal: abortSignal,
       },
     });
     if (serverConfigData.serverConfig.features.includes(ServerFeature.OAuth)) {
-      const oauthProvidersData = await this.gqlService.gql({
+      const oauthProvidersData = await gql({
         query: oauthProvidersQuery,
         context: {
           signal: abortSignal,

packages/frontend/core/src/modules/cloud/stores/server-list.ts

@@ -31,11 +31,17 @@ export class ServerListStore extends Store {
   }
 
   addServer(server: ServerMetadata, serverConfig: ServerConfig) {
-    this.updateServerConfig(server.id, serverConfig);
     const oldServers =
       this.globalStateService.globalState.get<ServerMetadata[]>('serverList') ??
       [];
 
+    if (oldServers.some(s => s.baseUrl === server.baseUrl)) {
+      throw new Error(
+        'Server with same base url already exists, ' + server.baseUrl
+      );
+    }
+
+    this.updateServerConfig(server.id, serverConfig);
     this.globalStateService.globalState.set<ServerMetadata[]>('serverList', [
       ...oldServers,
       server,