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,

packages/frontend/core/src/modules/desktop-api/service/desktop-api.ts

@@ -15,7 +15,7 @@ import {
   useNavigationType,
 } from 'react-router-dom';
 
-import { AuthService, ServersService } from '../../cloud';
+import { AuthService, DefaultServerService, ServersService } from '../../cloud';
 import type { DesktopApi } from '../entities/electron-api';
 
 @OnEvent(ApplicationStarted, e => e.setupStartListener)
@@ -138,20 +138,26 @@ export class DesktopApiService extends Service {
   }
 
   private setupAuthRequestEvent() {
-    this.events.ui.onAuthenticationRequest(({ method, payload }) => {
+    this.events.ui.onAuthenticationRequest(({ method, payload, server }) => {
       (async () => {
         if (!(await this.api.handler.ui.isActiveTab())) {
           return;
         }
 
-        // TODO: support multiple servers
-        const affineCloudServer = this.framework
-          .get(ServersService)
-          .server$('affine-cloud').value;
-        if (!affineCloudServer) {
+        // Dynamically get these services to avoid circular dependencies
+        const serversService = this.framework.get(ServersService);
+        const defaultServerService = this.framework.get(DefaultServerService);
+
+        let targetServer;
+        if (server) {
+          targetServer = await serversService.addOrGetServerByBaseUrl(server);
+        } else {
+          targetServer = defaultServerService.server;
+        }
+        if (!targetServer) {
           throw new Error('Affine Cloud server not found');
         }
-        const authService = affineCloudServer.scope.get(AuthService);
+        const authService = targetServer.scope.get(AuthService);
 
         switch (method) {
           case 'magic-link': {

packages/frontend/core/src/modules/dialogs/constant.ts

@@ -31,6 +31,9 @@ export type GLOBAL_DIALOG_SCHEMA = {
     workspaceMetadata?: WorkspaceMetadata | null;
     scrollAnchor?: string;
   }) => void;
+  'sign-in': (props: { server?: string; step?: 'sign-in' }) => void;
+  'change-password': (props: { server?: string }) => void;
+  'verify-email': (props: { server?: string; changeEmail?: boolean }) => void;
 };
 
 export type WORKSPACE_DIALOG_SCHEMA = {

packages/frontend/core/src/modules/favorite/stores/favorite.ts

@@ -1,4 +1,3 @@
-import { WorkspaceFlavour } from '@affine/env/workspace';
 import type { WorkspaceDBService, WorkspaceService } from '@toeverything/infra';
 import { LiveData, Store } from '@toeverything/infra';
 import { map } from 'rxjs';
@@ -27,7 +26,7 @@ export class FavoriteStore extends Store {
     // if is local workspace or no account, use __local__ userdata
     // sometimes we may have cloud workspace but no account for a short time, we also use __local__ userdata
     if (
-      this.workspaceService.workspace.meta.flavour === WorkspaceFlavour.LOCAL ||
+      this.workspaceService.workspace.meta.flavour === 'local' ||
       !this.authService
     ) {
       return new LiveData(this.workspaceDBService.userdataDB('__local__'));

packages/frontend/core/src/modules/import-template/services/import.ts

@@ -1,4 +1,3 @@
-import type { WorkspaceFlavour } from '@affine/env/workspace';
 import { type DocMode, ZipTransformer } from '@blocksuite/affine/blocks';
 import type { WorkspaceMetadata, WorkspacesService } from '@toeverything/infra';
 import { DocsService, Service } from '@toeverything/infra';
@@ -36,7 +35,7 @@ export class ImportTemplateService extends Service {
   }
 
   async importToNewWorkspace(
-    flavour: WorkspaceFlavour,
+    flavour: string,
     workspaceName: string,
     docBinary: Uint8Array
     // todo: support doc mode on init

packages/frontend/core/src/modules/permissions/entities/permission.ts

@@ -1,5 +1,4 @@
 import { DebugLogger } from '@affine/debug';
-import { WorkspaceFlavour } from '@affine/env/workspace';
 import type { WorkspaceService } from '@toeverything/infra';
 import {
   backoffRetry,
@@ -34,10 +33,7 @@ export class WorkspacePermission extends Entity {
   revalidate = effect(
     exhaustMap(() => {
       return fromPromise(async signal => {
-        if (
-          this.workspaceService.workspace.flavour ===
-          WorkspaceFlavour.AFFINE_CLOUD
-        ) {
+        if (this.workspaceService.workspace.flavour !== 'local') {
           return await this.store.fetchIsOwner(
             this.workspaceService.workspace.id,
             signal

packages/frontend/core/src/modules/share-doc/services/share-docs-list.ts

@@ -1,4 +1,3 @@
-import { WorkspaceFlavour } from '@affine/env/workspace';
 import type { WorkspaceService } from '@toeverything/infra';
 import { Service } from '@toeverything/infra';
 
@@ -10,7 +9,7 @@ export class ShareDocsListService extends Service {
   }
 
   shareDocs =
-    this.workspaceService.workspace.flavour === WorkspaceFlavour.AFFINE_CLOUD
+    this.workspaceService.workspace.flavour !== 'local'
       ? this.framework.createEntity(ShareDocsList)
       : null;
 }

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

@@ -1,8 +1,8 @@
 import { type Framework, GlobalContextService } from '@toeverything/infra';
 
-import { ServersService } from '../cloud/services/servers';
+import { ServersService } from '../cloud';
 import { TelemetryService } from './services/telemetry';
 
 export function configureTelemetryModule(framework: Framework) {
-  framework.service(TelemetryService, [ServersService, GlobalContextService]);
+  framework.service(TelemetryService, [GlobalContextService, ServersService]);
 }

packages/frontend/core/src/modules/telemetry/services/telemetry.ts

@@ -1,50 +1,58 @@
 import { mixpanel } from '@affine/track';
 import type { GlobalContextService } from '@toeverything/infra';
-import { ApplicationStarted, OnEvent, Service } from '@toeverything/infra';
+import {
+  ApplicationStarted,
+  LiveData,
+  OnEvent,
+  Service,
+} from '@toeverything/infra';
 
-import { AccountChanged, type AuthAccountInfo, AuthService } from '../../cloud';
-import { AccountLoggedOut } from '../../cloud/services/auth';
-import type { ServersService } from '../../cloud/services/servers';
+import type { AuthAccountInfo, Server, ServersService } from '../../cloud';
 
 @OnEvent(ApplicationStarted, e => e.onApplicationStart)
-@OnEvent(AccountChanged, e => e.updateIdentity)
-@OnEvent(AccountLoggedOut, e => e.onAccountLoggedOut)
 export class TelemetryService extends Service {
-  private readonly authService;
+  private readonly disposableFns: (() => void)[] = [];
+
+  private readonly currentAccount$ =
+    this.globalContextService.globalContext.serverId.$.selector(id =>
+      id
+        ? this.serversService.server$(id)
+        : new LiveData<Server | undefined>(undefined)
+    )
+      .flat()
+      .selector(server => server?.account$)
+      .flat();
+
   constructor(
-    serversService: ServersService,
-    private readonly globalContextService: GlobalContextService
+    private readonly globalContextService: GlobalContextService,
+    private readonly serversService: ServersService
   ) {
     super();
 
     // TODO: support multiple servers
-    const affineCloudServer = serversService.server$('affine-cloud').value;
-    if (!affineCloudServer) {
-      throw new Error('affine-cloud server not found');
-    }
-    this.authService = affineCloudServer.scope.get(AuthService);
-  }
 
-  onApplicationStart() {
-    const account = this.authService.session.account$.value;
-    this.updateIdentity(account);
-    this.registerMiddlewares();
-  }
-
-  updateIdentity(account: AuthAccountInfo | null) {
-    if (!account) {
-      return;
-    }
-    mixpanel.identify(account.id);
-    mixpanel.people.set({
-      $email: account.email,
-      $name: account.label,
-      $avatar: account.avatar,
+    let prevAccount: AuthAccountInfo | null = null;
+    const unsubscribe = this.currentAccount$.subscribe(account => {
+      if (prevAccount) {
+        mixpanel.reset();
+      }
+      prevAccount = account ?? null;
+      if (account) {
+        mixpanel.identify(account.id);
+        mixpanel.people.set({
+          $email: account.email,
+          $name: account.label,
+          $avatar: account.avatar,
+        });
+      }
+    });
+    this.disposableFns.push(() => {
+      unsubscribe.unsubscribe();
     });
   }
 
-  onAccountLoggedOut() {
-    mixpanel.reset();
+  onApplicationStart() {
+    this.registerMiddlewares();
   }
 
   registerMiddlewares() {
@@ -59,7 +67,7 @@ export class TelemetryService extends Service {
     );
   }
 
-  extractGlobalContext(): { page?: string } {
+  extractGlobalContext(): { page?: string; serverId?: string } {
     const globalContext = this.globalContextService.globalContext;
     const page = globalContext.isDoc.get()
       ? globalContext.isTrashDoc.get()
@@ -76,11 +84,12 @@ export class TelemetryService extends Service {
             : globalContext.isTag.get()
               ? 'tag'
               : undefined;
-    return { page };
+    const serverId = globalContext.serverId.get() ?? undefined;
+    return { page, serverId };
   }
 
   override dispose(): void {
-    this.disposables.forEach(dispose => dispose());
+    this.disposableFns.forEach(dispose => dispose());
     super.dispose();
   }
 }

packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts

@@ -1,5 +1,4 @@
 import { DebugLogger } from '@affine/debug';
-import { WorkspaceFlavour } from '@affine/env/workspace';
 import {
   createWorkspaceMutation,
   deleteWorkspaceMutation,
@@ -8,7 +7,6 @@ import {
 } from '@affine/graphql';
 import { DocCollection } from '@blocksuite/affine/store';
 import {
-  ApplicationStarted,
   type BlobStorage,
   catchErrorInto,
   type DocStorage,
@@ -16,22 +14,24 @@ import {
   fromPromise,
   type GlobalState,
   LiveData,
+  ObjectPool,
   onComplete,
-  OnEvent,
   onStart,
+  Service,
   type Workspace,
   type WorkspaceEngineProvider,
   type WorkspaceFlavourProvider,
+  type WorkspaceFlavoursProvider,
   type WorkspaceMetadata,
   type WorkspaceProfileInfo,
 } from '@toeverything/infra';
-import { effect, getAFFiNEWorkspaceSchema, Service } from '@toeverything/infra';
+import { effect, getAFFiNEWorkspaceSchema } from '@toeverything/infra';
 import { isEqual } from 'lodash-es';
 import { nanoid } from 'nanoid';
-import { EMPTY, map, mergeMap } from 'rxjs';
+import { EMPTY, map, mergeMap, Observable, switchMap } from 'rxjs';
 import { applyUpdate, encodeStateAsUpdate } from 'yjs';
 
-import type { Server } from '../../cloud';
+import type { Server, ServersService } from '../../cloud';
 import {
   AccountChanged,
   AuthService,
@@ -40,7 +40,6 @@ import {
   WebSocketService,
   WorkspaceServerService,
 } from '../../cloud';
-import type { ServersService } from '../../cloud/services/servers';
 import type { WorkspaceEngineStorageProvider } from '../providers/engine';
 import { BroadcastChannelAwarenessConnection } from './engine/awareness-broadcast-channel';
 import { CloudAwarenessConnection } from './engine/awareness-cloud';
@@ -49,41 +48,41 @@ import { StaticBlobStorage } from './engine/blob-static';
 import { CloudDocEngineServer } from './engine/doc-cloud';
 import { CloudStaticDocStorage } from './engine/doc-cloud-static';
 
-const CLOUD_WORKSPACES_CACHE_KEY = 'cloud-workspace:';
+const getCloudWorkspaceCacheKey = (serverId: string) => {
+  if (serverId === 'affine-cloud') {
+    return 'cloud-workspace:'; // FOR BACKWARD COMPATIBILITY
+  }
+  return `selfhosted-workspace-${serverId}:`;
+};
 
 const logger = new DebugLogger('affine:cloud-workspace-flavour-provider');
 
-@OnEvent(ApplicationStarted, e => e.revalidate)
-@OnEvent(AccountChanged, e => e.revalidate)
-export class CloudWorkspaceFlavourProviderService
-  extends Service
-  implements WorkspaceFlavourProvider
-{
+class CloudWorkspaceFlavourProvider implements WorkspaceFlavourProvider {
   private readonly authService: AuthService;
   private readonly webSocketService: WebSocketService;
   private readonly fetchService: FetchService;
   private readonly graphqlService: GraphQLService;
-  private readonly affineCloudServer: Server;
+
+  private readonly unsubscribeAccountChanged: () => void;
 
   constructor(
     private readonly globalState: GlobalState,
     private readonly storageProvider: WorkspaceEngineStorageProvider,
-    serversService: ServersService
+    private readonly server: Server
   ) {
-    super();
-    // TODO: support multiple servers
-    const affineCloudServer = serversService.server$('affine-cloud').value;
-    if (!affineCloudServer) {
-      throw new Error('affine-cloud server not found');
-    }
-    this.affineCloudServer = affineCloudServer;
-    this.authService = affineCloudServer.scope.get(AuthService);
-    this.webSocketService = affineCloudServer.scope.get(WebSocketService);
-    this.fetchService = affineCloudServer.scope.get(FetchService);
-    this.graphqlService = affineCloudServer.scope.get(GraphQLService);
+    this.authService = server.scope.get(AuthService);
+    this.webSocketService = server.scope.get(WebSocketService);
+    this.fetchService = server.scope.get(FetchService);
+    this.graphqlService = server.scope.get(GraphQLService);
+    this.unsubscribeAccountChanged = this.server.scope.eventBus.on(
+      AccountChanged,
+      () => {
+        this.revalidate();
+      }
+    );
   }
 
-  flavour: WorkspaceFlavour = WorkspaceFlavour.AFFINE_CLOUD;
+  flavour = this.server.id;
 
   async deleteWorkspace(id: string): Promise<void> {
     await this.graphqlService.gql({
@@ -95,6 +94,7 @@ export class CloudWorkspaceFlavourProviderService
     this.revalidate();
     await this.waitForLoaded();
   }
+
   async createWorkspace(
     initial: (
       docCollection: DocCollection,
@@ -139,7 +139,7 @@ export class CloudWorkspaceFlavourProviderService
 
     return {
       id: workspaceId,
-      flavour: WorkspaceFlavour.AFFINE_CLOUD,
+      flavour: this.server.id,
     };
   }
   revalidate = effect(
@@ -169,7 +169,7 @@ export class CloudWorkspaceFlavourProviderService
             accountId,
             workspaces: ids.map(({ id, initialized }) => ({
               id,
-              flavour: WorkspaceFlavour.AFFINE_CLOUD,
+              flavour: this.server.id,
               initialized,
             })),
           };
@@ -181,7 +181,7 @@ export class CloudWorkspaceFlavourProviderService
                 return a.id.localeCompare(b.id);
               });
               this.globalState.set(
-                CLOUD_WORKSPACES_CACHE_KEY + accountId,
+                getCloudWorkspaceCacheKey(this.server.id) + accountId,
                 sorted
               );
               if (!isEqual(this.workspaces$.value, sorted)) {
@@ -202,7 +202,9 @@ export class CloudWorkspaceFlavourProviderService
       ({ accountId }) => {
         if (accountId) {
           this.workspaces$.next(
-            this.globalState.get(CLOUD_WORKSPACES_CACHE_KEY + accountId) ?? []
+            this.globalState.get(
+              getCloudWorkspaceCacheKey(this.server.id) + accountId
+            ) ?? []
           );
         } else {
           this.workspaces$.next([]);
@@ -295,9 +297,7 @@ export class CloudWorkspaceFlavourProviderService
 
   onWorkspaceInitialized(workspace: Workspace): void {
     // bind the workspace to the affine cloud server
-    workspace.scope
-      .get(WorkspaceServerService)
-      .bindServer(this.affineCloudServer);
+    workspace.scope.get(WorkspaceServerService).bindServer(this.server);
   }
 
   private async getIsOwner(workspaceId: string, signal?: AbortSignal) {
@@ -315,4 +315,61 @@ export class CloudWorkspaceFlavourProviderService
   private waitForLoaded() {
     return this.isRevalidating$.waitFor(loading => !loading);
   }
+
+  dispose() {
+    this.revalidate.unsubscribe();
+    this.unsubscribeAccountChanged();
+  }
+}
+
+export class CloudWorkspaceFlavoursProvider
+  extends Service
+  implements WorkspaceFlavoursProvider
+{
+  constructor(
+    private readonly globalState: GlobalState,
+    private readonly storageProvider: WorkspaceEngineStorageProvider,
+    private readonly serversService: ServersService
+  ) {
+    super();
+  }
+
+  workspaceFlavours$ = LiveData.from<WorkspaceFlavourProvider[]>(
+    this.serversService.servers$.pipe(
+      switchMap(servers => {
+        const refs = servers.map(server => {
+          const exists = this.pool.get(server.id);
+          if (exists) {
+            return exists;
+          }
+          const provider = new CloudWorkspaceFlavourProvider(
+            this.globalState,
+            this.storageProvider,
+            server
+          );
+          provider.revalidate();
+          const ref = this.pool.put(server.id, provider);
+          return ref;
+        });
+
+        return new Observable<WorkspaceFlavourProvider[]>(subscribe => {
+          subscribe.next(refs.map(ref => ref.obj));
+          return () => {
+            refs.forEach(ref => {
+              ref.release();
+            });
+          };
+        });
+      })
+    ),
+    [] as any
+  );
+
+  private readonly pool = new ObjectPool<string, CloudWorkspaceFlavourProvider>(
+    {
+      onDelete(obj) {
+        obj.dispose();
+      },
+    }
+  );
 }

packages/frontend/core/src/modules/workspace-engine/impls/local.ts

@@ -1,11 +1,12 @@
 import { DebugLogger } from '@affine/debug';
-import { WorkspaceFlavour } from '@affine/env/workspace';
 import { DocCollection } from '@blocksuite/affine/store';
 import type {
   BlobStorage,
   DocStorage,
+  FrameworkProvider,
   WorkspaceEngineProvider,
   WorkspaceFlavourProvider,
+  WorkspaceFlavoursProvider,
   WorkspaceMetadata,
   WorkspaceProfileInfo,
 } from '@toeverything/infra';
@@ -54,17 +55,13 @@ export function setLocalWorkspaceIds(
   );
 }
 
-export class LocalWorkspaceFlavourProvider
-  extends Service
-  implements WorkspaceFlavourProvider
-{
+class LocalWorkspaceFlavourProvider implements WorkspaceFlavourProvider {
   constructor(
-    private readonly storageProvider: WorkspaceEngineStorageProvider
-  ) {
-    super();
-  }
+    private readonly storageProvider: WorkspaceEngineStorageProvider,
+    private readonly framework: FrameworkProvider
+  ) {}
 
-  flavour: WorkspaceFlavour = WorkspaceFlavour.LOCAL;
+  flavour = 'local';
   notifyChannel = new BroadcastChannel(
     LOCAL_WORKSPACE_CHANGED_BROADCAST_CHANNEL_KEY
   );
@@ -72,9 +69,8 @@ export class LocalWorkspaceFlavourProvider
   async deleteWorkspace(id: string): Promise<void> {
     setLocalWorkspaceIds(ids => ids.filter(x => x !== id));
 
-    const electronApi = this.framework.getOptional(DesktopApiService);
-
-    if (BUILD_CONFIG.isElectron && electronApi) {
+    if (BUILD_CONFIG.isElectron) {
+      const electronApi = this.framework.get(DesktopApiService);
       await electronApi.handler.workspace.delete(id);
     }
 
@@ -116,7 +112,7 @@ export class LocalWorkspaceFlavourProvider
     // notify all browser tabs, so they can update their workspace list
     this.notifyChannel.postMessage(id);
 
-    return { id, flavour: WorkspaceFlavour.LOCAL };
+    return { id, flavour: 'local' };
   }
   workspaces$ = LiveData.from(
     new Observable<WorkspaceMetadata[]>(subscriber => {
@@ -124,7 +120,7 @@ export class LocalWorkspaceFlavourProvider
       const emit = () => {
         const value = getLocalWorkspaceIds().map(id => ({
           id,
-          flavour: WorkspaceFlavour.LOCAL,
+          flavour: 'local',
         }));
         if (isEqual(last, value)) return;
         subscriber.next(value);
@@ -199,3 +195,18 @@ export class LocalWorkspaceFlavourProvider
     };
   }
 }
+
+export class LocalWorkspaceFlavoursProvider
+  extends Service
+  implements WorkspaceFlavoursProvider
+{
+  constructor(
+    private readonly storageProvider: WorkspaceEngineStorageProvider
+  ) {
+    super();
+  }
+
+  workspaceFlavours$ = new LiveData<WorkspaceFlavourProvider[]>([
+    new LocalWorkspaceFlavourProvider(this.storageProvider, this.framework),
+  ]);
+}

packages/frontend/core/src/modules/workspace-engine/index.ts

@@ -1,19 +1,19 @@
 import {
   type Framework,
   GlobalState,
-  WorkspaceFlavourProvider,
+  WorkspaceFlavoursProvider,
 } from '@toeverything/infra';
 
 import { ServersService } from '../cloud/services/servers';
 import { DesktopApiService } from '../desktop-api';
-import { CloudWorkspaceFlavourProviderService } from './impls/cloud';
+import { CloudWorkspaceFlavoursProvider } from './impls/cloud';
 import { IndexedDBBlobStorage } from './impls/engine/blob-indexeddb';
 import { SqliteBlobStorage } from './impls/engine/blob-sqlite';
 import { IndexedDBDocStorage } from './impls/engine/doc-indexeddb';
 import { SqliteDocStorage } from './impls/engine/doc-sqlite';
 import {
   LOCAL_WORKSPACE_LOCAL_STORAGE_KEY,
-  LocalWorkspaceFlavourProvider,
+  LocalWorkspaceFlavoursProvider,
 } from './impls/local';
 import { WorkspaceEngineStorageProvider } from './providers/engine';
 
@@ -21,17 +21,14 @@ export { CloudBlobStorage } from './impls/engine/blob-cloud';
 
 export function configureBrowserWorkspaceFlavours(framework: Framework) {
   framework
-    .impl(WorkspaceFlavourProvider('LOCAL'), LocalWorkspaceFlavourProvider, [
+    .impl(WorkspaceFlavoursProvider('LOCAL'), LocalWorkspaceFlavoursProvider, [
       WorkspaceEngineStorageProvider,
     ])
-    .service(CloudWorkspaceFlavourProviderService, [
+    .impl(WorkspaceFlavoursProvider('CLOUD'), CloudWorkspaceFlavoursProvider, [
       GlobalState,
       WorkspaceEngineStorageProvider,
       ServersService,
-    ])
-    .impl(WorkspaceFlavourProvider('CLOUD'), p =>
-      p.get(CloudWorkspaceFlavourProviderService)
-    );
+    ]);
 }
 
 export function configureIndexedDBWorkspaceEngineStorageProvider(