Repository: novuhq/novu
File: apps/ws/src/socket/ws.gateway.ts (TypeScript)

Code before:

  private disconnect(socket: Socket) {
    socket.disconnect();
  }

  async gracefulShutdown(): Promise<void> {
    try {
      if (!this.server) {
        Logger.error('WS server was not initialized while executing shutdown', LOG_CONTEXT);

        return;
      }

      Logger.log('Closing WS server for incoming new connections', LOG_CONTEXT);
      this.server.close();

      Logger.log('Disconnecting active sockets connections', LOG_CONTEXT);
      this.server.sockets.disconnectSockets();
    } catch (e) {
      Logger.error(e, 'Unexpected exception was thrown while graceful shut down was executed', LOG_CONTEXT);
      throw e;
    } finally {
      Logger.log(`Graceful shutdown down has finished`, LOG_CONTEXT);
    }
  }

  async onModuleDestroy(): Promise<void> {
    this.isShutdown = true;
    await this.gracefulShutdown();
  }

  afterInit() {
    if (!!process.env.SOCKET_IO_ADMIN_USERNAME && !!process.env.SOCKET_IO_ADMIN_PASSWORD_HASH) {
      // For more information on how to use the admin UI, see https://socket.io/docs/v4/admin-ui/
      instrument(this.server, {
        auth: {
          type: 'basic',
          username: process.env.SOCKET_IO_ADMIN_USERNAME,
          password: process.env.SOCKET_IO_ADMIN_PASSWORD_HASH,
        },
        mode: process.env.NODE_ENV === 'production' ? 'production' : 'development',
        namespaceName: '/admin',
      });
    }
  }
}


Review comment (suggestion):
maybe we should send the `contextKey` in the event data to allow doing the custom logic for those using `@novu/js` package? 🤔 

Code after:
      );
    }

    Logger.log(`Connection request accepted for ${subscriber._id}`, LOG_CONTEXT);

    await this.subscriberOnlineService.handleConnection(subscriber);
  }

  async sendMessage(userId: string, event: string, data: any, contextKeys?: string[]) {
    if (!this.server) {
      Logger.error('No sw server available to send message', LOG_CONTEXT);

      return;
    }

    const messageContextKeys = contextKeys || [];

    if (messageContextKeys.length === 0) {
      // Send to default room (connections without context)
      Logger.log(`Sending event ${event} to ${userId}:${DEFAULT_CONTEXT_ROOM} (no context)`, LOG_CONTEXT);
      this.server.to(`${userId}:${DEFAULT_CONTEXT_ROOM}`).emit(event, data);
    } else {
      // Send to all matching context rooms
      Logger.log(`Sending event ${event} to ${userId} with contexts: ${messageContextKeys.join(', ')}`, LOG_CONTEXT);
      for (const contextKey of messageContextKeys) {
        this.server.to(`${userId}:${contextKey}`).emit(event, data);
      }
    }
  }

  private disconnect(socket: Socket) {
    socket.disconnect();
  }

  async gracefulShutdown(): Promise<void> {
    try {
      if (!this.server) {
        Logger.error('WS server was not initialized while executing shutdown', LOG_CONTEXT);

        return;
      }

      Logger.log('Closing WS server for incoming new connections', LOG_CONTEXT);
      this.server.close();

      Logger.log('Disconnecting active sockets connections', LOG_CONTEXT);
      this.server.sockets.disconnectSockets();
    } catch (e) {
      Logger.error(e, 'Unexpected exception was thrown while graceful shut down was executed', LOG_CONTEXT);
      throw e;
    } finally {