gateway

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