fix(core): disconnect ws when user logout (#8188)

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

@@ -59,7 +59,7 @@ export function configureCloudModule(framework: Framework) {
   framework
     .service(FetchService)
     .service(GraphQLService, [FetchService])
-    .service(WebSocketService)
+    .service(WebSocketService, [AuthService])
     .service(ServerConfigService)
     .entity(ServerConfig, [ServerConfigStore])
     .store(ServerConfigStore, [GraphQLService])

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

@@ -1,37 +1,47 @@
-import { OnEvent, Service } from '@toeverything/infra';
-import type { Socket } from 'socket.io-client';
+import { ApplicationStarted, OnEvent, Service } from '@toeverything/infra';
 import { Manager } from 'socket.io-client';
 
 import { getAffineCloudBaseUrl } from '../services/fetch';
+import type { AuthService } from './auth';
 import { AccountChanged } from './auth';
 
-@OnEvent(AccountChanged, e => e.reconnect)
+@OnEvent(AccountChanged, e => e.update)
+@OnEvent(ApplicationStarted, e => e.update)
 export class WebSocketService extends Service {
   ioManager: Manager = new Manager(`${getAffineCloudBaseUrl()}/`, {
     autoConnect: false,
     transports: ['websocket'],
     secure: location.protocol === 'https:',
   });
-  sockets: Set<Socket> = new Set();
+  socket = this.ioManager.socket('/');
+  refCount = 0;
 
-  constructor() {
+  constructor(private readonly authService: AuthService) {
     super();
   }
 
-  newSocket(): Socket {
-    const socket = this.ioManager.socket('/');
-    this.sockets.add(socket);
-
-    return socket;
+  /**
+   * Connect socket, with automatic connect and reconnect logic.
+   * External code should not call `socket.connect()` or `socket.disconnect()` manually.
+   * When socket is no longer needed, call `dispose()` to clean up resources.
+   */
+  connect() {
+    this.refCount++;
+    this.update();
+    return {
+      socket: this.socket,
+      dispose: () => {
+        this.refCount--;
+        this.update();
+      },
+    };
   }
 
-  reconnect(): void {
-    for (const socket of this.sockets) {
-      socket.disconnect();
-    }
-
-    for (const socket of this.sockets) {
-      socket.connect();
+  update(): void {
+    if (this.authService.session.account$.value && this.refCount > 0) {
+      this.socket.connect();
+    } else {
+      this.socket.disconnect();
     }
   }
 }