gateway

src/IOC/ioc.config.ts

@@ -58,6 +58,7 @@ import { ChatbotService } from "../modules/chatbot/providers/chatbot.service";
 import { DataContextService } from "../modules/chatbot/providers/data-context.service";
 import { LLMService } from "../modules/chatbot/providers/llm.service";
 import { WebSocketAuthService } from "../modules/chatbot/providers/websocket-auth.service";
+import { ChatbotGateway } from "../modules/chatbot/chatbot.gateway";
 import { ChatSessionRepository, createChatSessionRepository } from "../modules/chatbot/repositories/chat-session.repository";
 import { ChatMessageRepository, createChatMessageRepository } from "../modules/chatbot/repositories/chat-message.repository";
 import { ContactUsRepo, CreateContactUsRepo } from "../modules/contact-us/contactUs.repository";
@@ -207,6 +208,7 @@ const containerModules = new AsyncContainerModule(async (bind) => {
   bind<LLMService>(IOCTYPES.ChatbotLLMService).to(LLMService).inSingletonScope();
   bind<DataContextService>(IOCTYPES.ChatbotDataContextService).to(DataContextService).inSingletonScope();
   bind<WebSocketAuthService>(IOCTYPES.ChatbotWebSocketAuthService).to(WebSocketAuthService).inSingletonScope();
+  bind<ChatbotGateway>(IOCTYPES.ChatbotGateway).to(ChatbotGateway).inSingletonScope();
   // #endregion
   // #region repository
   bind<CategoryRepository>(IOCTYPES.CategoryRepository).toDynamicValue(createCategoryRepo).inSingletonScope();

src/modules/chatbot/chatbot.gateway.ts

@@ -0,0 +1,330 @@
+import { Server as HttpServer } from "http";
+
+import { inject, injectable } from "inversify";
+import { Server, Socket } from "socket.io";
+
+import { ChatbotService } from "./providers/chatbot.service";
+import { WebSocketAuthService } from "./providers/websocket-auth.service";
+import { SendMessageDto } from "./DTO/send-message.dto";
+import { WEBSOCKET_EVENTS } from "./constants/chatbot.constants";
+import { IOCTYPES } from "../../IOC/ioc.types";
+import { Logger } from "../../core/logging/logger";
+import { AuthenticatedSocket, WebSocketResponse } from "./interfaces/websocket.interface";
+import * as ulidLib from "ulid";
+
+@injectable()
+export class ChatbotGateway {
+  private io: Server;
+  private readonly logger: Logger;
+  private readonly CHATBOT_ULID_COOKIE_NAME = "chatbot_session_id";
+
+  constructor(
+    @inject(IOCTYPES.ChatbotService) private chatbotService: ChatbotService,
+    @inject(IOCTYPES.ChatbotWebSocketAuthService) private wsAuthService: WebSocketAuthService,
+  ) {
+    this.logger = new Logger("ChatbotGateway");
+  }
+
+  public initialize(server: HttpServer): Server {
+    this.io = new Server(server, {
+      path: "/ws-chatbot",
+      cors: {
+        origin: true,
+        allowedHeaders: ["Authorization"],
+        credentials: true,
+        methods: ["GET", "POST"],
+      },
+    });
+
+    this.io.on("connection", (socket: Socket) => this.handleConnection(socket));
+    this.logger.info("ChatbotGateway initialized on /ws-chatbot");
+    return this.io;
+  }
+
+  private async handleConnection(socket: Socket): Promise<void> {
+    try {
+      // Try to authenticate if token is provided
+      const token = socket.handshake.query.token as string;
+      let ulid: string | undefined;
+
+      if (token) {
+        // Authenticated connection
+        const authResult = await this.wsAuthService.authenticateClient(socket as AuthenticatedSocket);
+        if (!authResult.success) {
+          this.wsAuthService.handleAuthenticationFailure(socket, authResult.error || "Authentication failed");
+          return;
+        }
+        this.wsAuthService.emitAuthenticationSuccess(socket, authResult.user!);
+        ulid = authResult.user!.id;
+      } else {
+        // Anonymous connection - generate or get ULID from cookie
+        const cookieHeader = socket.handshake.headers.cookie;
+        ulid = this.extractUlidFromCookie(cookieHeader);
+        if (!ulid) {
+          ulid = ulidLib.ulid();
+        }
+        socket.data = { user: { id: ulid, sub: ulid } };
+      }
+
+      this.logger.info(`Client connected: ${socket.id}, ULID: ${ulid}`);
+
+      // Register event handlers
+      this.registerEventHandlers(socket, ulid);
+    } catch (error) {
+      this.logger.error("Connection error", error);
+      socket.emit(WEBSOCKET_EVENTS.ERROR, {
+        status: "error",
+        message: "Connection failed",
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+      socket.disconnect();
+    }
+  }
+
+  private extractUlidFromCookie(cookieHeader?: string | string[]): string | undefined {
+    if (!cookieHeader || typeof cookieHeader !== "string") {
+      return undefined;
+    }
+
+    const cookies = cookieHeader.split(";").reduce(
+      (acc, cookie) => {
+        const [key, value] = cookie.trim().split("=");
+        if (key && value) {
+          acc[key] = decodeURIComponent(value);
+        }
+        return acc;
+      },
+      {} as Record<string, string>,
+    );
+
+    return cookies[this.CHATBOT_ULID_COOKIE_NAME];
+  }
+
+  private registerEventHandlers(socket: Socket, ulid: string): void {
+    // Session management
+    socket.on(WEBSOCKET_EVENTS.CREATE_SESSION, () => this.handleCreateSession(socket, ulid));
+    socket.on(WEBSOCKET_EVENTS.JOIN_CHAT, (sessionId: string) => this.handleJoinChat(socket, ulid, sessionId));
+    socket.on(WEBSOCKET_EVENTS.LEAVE_CHAT, (sessionId: string) => this.handleLeaveChat(socket, sessionId));
+
+    // Message handling
+    socket.on(WEBSOCKET_EVENTS.SEND_MESSAGE, (data: SendMessageDto) => this.handleSendMessage(socket, ulid, data));
+    socket.on(WEBSOCKET_EVENTS.SEND_MESSAGE + "_stream", (data: SendMessageDto) => this.handleSendMessageStream(socket, ulid, data));
+
+    // Typing indicators
+    socket.on(WEBSOCKET_EVENTS.TYPING_START, (sessionId: string) => this.handleTypingStart(socket, sessionId));
+    socket.on(WEBSOCKET_EVENTS.TYPING_STOP, (sessionId: string) => this.handleTypingStop(socket, sessionId));
+
+    // Disconnect
+    socket.on(WEBSOCKET_EVENTS.DISCONNECT, () => this.handleDisconnect(socket));
+  }
+
+  private async handleCreateSession(socket: Socket, ulid: string): Promise<void> {
+    try {
+      const session = await this.chatbotService.createChatSession(ulid);
+      socket.emit(WEBSOCKET_EVENTS.SESSION_CREATED, {
+        status: "success",
+        data: session,
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+    } catch (error) {
+      this.logger.error("Error creating session", error);
+      socket.emit(WEBSOCKET_EVENTS.SESSION_ERROR, {
+        status: "error",
+        message: "Failed to create session",
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+    }
+  }
+
+  private async handleJoinChat(socket: Socket, ulid: string, sessionId: string): Promise<void> {
+    try {
+      const session = await this.chatbotService.getChatSession(sessionId, ulid);
+      socket.join(`session_${sessionId}`);
+      (socket as AuthenticatedSocket).data.sessionId = sessionId;
+
+      socket.emit(WEBSOCKET_EVENTS.CHAT_JOINED, {
+        status: "success",
+        data: session,
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+
+      // Notify others in the session (if any)
+      socket.to(`session_${sessionId}`).emit(WEBSOCKET_EVENTS.USER_JOINED, {
+        status: "success",
+        data: { userId: ulid, sessionId },
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+    } catch (error) {
+      this.logger.error("Error joining chat", error);
+      socket.emit(WEBSOCKET_EVENTS.SESSION_ERROR, {
+        status: "error",
+        message: "Failed to join chat session",
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+    }
+  }
+
+  private handleLeaveChat(socket: Socket, sessionId: string): void {
+    socket.leave(`session_${sessionId}`);
+    delete (socket as AuthenticatedSocket).data.sessionId;
+
+    socket.emit(WEBSOCKET_EVENTS.CHAT_LEFT, {
+      status: "success",
+      message: "Left chat session",
+      timestamp: new Date().toISOString(),
+    } as WebSocketResponse);
+
+    socket.to(`session_${sessionId}`).emit(WEBSOCKET_EVENTS.USER_LEFT, {
+      status: "success",
+      data: { sessionId },
+      timestamp: new Date().toISOString(),
+    } as WebSocketResponse);
+  }
+
+  private async handleSendMessage(socket: Socket, ulid: string, data: SendMessageDto): Promise<void> {
+    try {
+      // Save user message
+      const userMessage = await this.chatbotService.sendMessage(ulid, data);
+
+      // Emit user message to client
+      socket.emit(WEBSOCKET_EVENTS.MESSAGE_RECEIVED, {
+        status: "success",
+        data: userMessage,
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+
+      // Emit to others in the session (if any)
+      socket.to(`session_${data.sessionId}`).emit(WEBSOCKET_EVENTS.MESSAGE_RECEIVED, {
+        status: "success",
+        data: userMessage,
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+
+      // Generate bot response asynchronously
+      this.generateBotResponseAsync(socket, data.sessionId, userMessage.id, ulid);
+    } catch (error) {
+      this.logger.error("Error sending message", error);
+      socket.emit(WEBSOCKET_EVENTS.MESSAGE_ERROR, {
+        status: "error",
+        message: "Failed to send message",
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+    }
+  }
+
+  private async handleSendMessageStream(socket: Socket, ulid: string, data: SendMessageDto): Promise<void> {
+    try {
+      const { userMessage, streamGenerator } = await this.chatbotService.sendMessageStream(ulid, data);
+
+      // Emit user message
+      socket.emit(WEBSOCKET_EVENTS.MESSAGE_RECEIVED, {
+        status: "success",
+        data: userMessage,
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+
+      // Emit to others in the session (if any)
+      socket.to(`session_${data.sessionId}`).emit(WEBSOCKET_EVENTS.MESSAGE_RECEIVED, {
+        status: "success",
+        data: userMessage,
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+
+      // Start streaming bot response
+      socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE_START, {
+        status: "success",
+        data: { userMessageId: userMessage.id },
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+
+      let fullResponse = "";
+      const stream = await streamGenerator();
+      for await (const chunk of stream) {
+        if (chunk) {
+          fullResponse += chunk;
+          socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE_CHUNK, {
+            status: "success",
+            data: { chunk, userMessageId: userMessage.id },
+            timestamp: new Date().toISOString(),
+          } as WebSocketResponse);
+        }
+      }
+
+      // Save bot response after streaming completes
+      if (fullResponse) {
+        await this.chatbotService.saveStreamedBotResponse(data.sessionId, userMessage.id, fullResponse, ulid);
+      }
+
+      // Emit final bot response
+      socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE_END, {
+        status: "success",
+        data: { userMessageId: userMessage.id, fullResponse },
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+    } catch (error) {
+      this.logger.error("Error sending stream message", error);
+      socket.emit(WEBSOCKET_EVENTS.MESSAGE_ERROR, {
+        status: "error",
+        message: "Failed to send message",
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+    }
+  }
+
+  private async generateBotResponseAsync(socket: Socket, sessionId: string, userMessageId: string, ulid: string): Promise<void> {
+    try {
+      // Get the session to check if it exists
+      const session = await this.chatbotService.getChatSession(sessionId, ulid);
+      if (!session) {
+        return;
+      }
+
+      // The bot response will be generated asynchronously by the service
+      // We'll need to emit it when it's ready
+      // For now, we'll rely on the service's internal async generation
+      // In a production system, you might want to use Redis pub/sub or similar
+      // to notify when the bot response is ready
+
+      // Emit that bot is typing
+      socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE_START, {
+        status: "success",
+        data: { userMessageId },
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+
+      // Note: The actual bot response will be generated by ChatbotService.generateBotResponse
+      // which runs asynchronously. In a real implementation, you might want to use
+      // Redis pub/sub or a message queue to notify when the response is ready
+    } catch (error) {
+      this.logger.error("Error generating bot response", error);
+      socket.emit(WEBSOCKET_EVENTS.BOT_RESPONSE, {
+        status: "error",
+        message: "Failed to generate bot response",
+        timestamp: new Date().toISOString(),
+      } as WebSocketResponse);
+    }
+  }
+
+  private handleTypingStart(socket: Socket, sessionId: string): void {
+    socket.to(`session_${sessionId}`).emit(WEBSOCKET_EVENTS.TYPING_START, {
+      status: "success",
+      data: { sessionId, userId: socket.data?.user?.id },
+      timestamp: new Date().toISOString(),
+    } as WebSocketResponse);
+  }
+
+  private handleTypingStop(socket: Socket, sessionId: string): void {
+    socket.to(`session_${sessionId}`).emit(WEBSOCKET_EVENTS.TYPING_STOP, {
+      status: "success",
+      data: { sessionId, userId: socket.data?.user?.id },
+      timestamp: new Date().toISOString(),
+    } as WebSocketResponse);
+  }
+
+  private handleDisconnect(socket: Socket): void {
+    const user = socket.data?.user;
+    if (user) {
+      this.logger.info(`Client disconnected: ${socket.id}, ULID: ${user.id}`);
+    }
+  }
+}

src/server.ts

@@ -11,6 +11,7 @@ import { Logger } from "./core/logging/logger";
 import { connectMongo } from "./db/connection";
 import { IOCTYPES } from "./IOC/ioc.types";
 import { ChatGateway } from "./modules/chat/chat.gateway";
+import { ChatbotGateway } from "./modules/chatbot/chatbot.gateway";
 import { StartWorker } from "./queues";
 import { gracefulShutdown } from "./utils/shutdown.utils";
 
@@ -28,6 +29,9 @@ async function bootStrap() {
     const chatGateway = container.get<ChatGateway>(IOCTYPES.ChatGateway);
     const wsServer = chatGateway.initialize(server);
 
+    const chatbotGateway = container.get<ChatbotGateway>(IOCTYPES.ChatbotGateway);
+    const chatbotWsServer = chatbotGateway.initialize(server);
+
     const worker = await StartWorker();
 
     //
@@ -36,9 +40,9 @@ async function bootStrap() {
       logger.info(`swagger documentation is serving on http://localhost:${PORT}/api-docs`);
     });
 
-    process.on("SIGINT", () => gracefulShutdown("SIGINT", server, wsServer, worker[0], logger));
+    process.on("SIGINT", () => gracefulShutdown("SIGINT", server, wsServer, worker[0], logger, chatbotWsServer));
 
-    process.on("SIGTERM", () => gracefulShutdown("SIGTERM", server, wsServer, worker[0], logger));
+    process.on("SIGTERM", () => gracefulShutdown("SIGTERM", server, wsServer, worker[0], logger, chatbotWsServer));
   } catch (error) {
     logger.error("Error starting the server.", error);
     process.exit(1);

src/utils/shutdown.utils.ts

@@ -5,16 +5,29 @@ import { Server as WSServer } from "socket.io";
 
 import { Logger } from "../core/logging/logger";
 
-export const gracefulShutdown = async (signal: string, server: Server, wsServer: WSServer, worker: Worker, logger: Logger) => {
+export const gracefulShutdown = async (
+  signal: string,
+  server: Server,
+  wsServer: WSServer,
+  worker: Worker,
+  logger: Logger,
+  chatbotWsServer?: WSServer,
+) => {
   try {
     logger.warn(`${signal} signal received: closing HTTP server and worker...`);
 
     logger.warn("close all socket connection..");
     wsServer.disconnectSockets();
+    if (chatbotWsServer) {
+      chatbotWsServer.disconnectSockets();
+    }
     logger.warn("all socket connection disconnected");
 
     logger.warn("Closing ws server...");
     wsServer.close();
+    if (chatbotWsServer) {
+      chatbotWsServer.close();
+    }
     logger.warn("ws server closed");
 
     server.close(async () => {