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

Code before:
    Logger.log(
      `Connection request received from ${subscriber._id} external id: ${subscriber.subscriberId} organization id: ${subscriber.organizationId}`,
      LOG_CONTEXT
    );

    await connection.join(subscriber._id);

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

    await this.subscriberOnlineService.handleConnection(subscriber);
  }

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

      return;
    }

    Logger.log(`Sending event ${event} message to ${userId}`, LOG_CONTEXT);

    this.server.to(userId).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 {
      Logger.log(`Graceful shutdown down has finished`, LOG_CONTEXT);
    }
  }

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

Review comment (performance):
To route events only to inbox sessions with matching contexts, there are two possible approaches:

1. **Room-based routing (chosen)**: Each inbox session joins a context-specific room on connection. Events are emitted directly to the appropriate rooms.

2. **JWT-based filtering**: Decode the JWT token for each connected client when sending an event, extract the contextKeys, and decide whether to emit.

I think the #1 approach is better for performance and scalability.

Code after:

    const contextKeys = subscriber.contextKeys || [];

    Logger.log(
      `Connection request received from ${subscriber._id} external id: ${subscriber.subscriberId} organization id: ${subscriber.organizationId}`,
      LOG_CONTEXT
    );

    await connection.join(subscriber._id);

    // Join context-specific rooms for efficient message routing
    if (contextKeys.length === 0) {
      // Join default room for connections without context
      await connection.join(`${subscriber._id}:${DEFAULT_CONTEXT_ROOM}`);
      Logger.log(`Connection ${connection.id} joined default room for ${subscriber._id}`, LOG_CONTEXT);
    } else {
      // Join a room for each context key
      for (const contextKey of contextKeys) {
        await connection.join(`${subscriber._id}:${contextKey}`);
      }
      Logger.log(
        `Connection ${connection.id} joined context rooms for ${subscriber._id}: ${contextKeys.join(', ')}`,
        LOG_CONTEXT
      );
    }

    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);
      }
    }