@xalia/agent 0.6.10 → 0.6.11

package/src/chat/server/connectionManager.ts

@@ -1,506 +0,0 @@
-import * as ws from "ws";
-import { v4 as uuidv4 } from "uuid";
-import { getLogger } from "@xalia/xmcp/sdk";
-
-import {
-  ServerConnectionReady,
-  ServerConnectionError,
-  ServerConnectionPing,
-  ClientConnectionMessage,
-  ServerConnectionMessage,
-} from "../protocol/connectionMessages";
-import { ChatErrorMessage, ChatFatalError } from "../protocol/errors";
-
-const logger = getLogger();
-
-const HEARTBEAT_INTERVAL_MS: number = parseInt(
-  process.env["HEARTBEAT_INTERVAL_MS"] || String(30 * 1000)
-);
-
-/**
- * Interface for sending messages to clients.
- */
-export interface IUserConnectionManager<ServerMsgT> {
-  /**
-   * Send message to all active connections of specific users.
-   * Handles user-to-connection routing internally.
-   */
-  sendToUsers(userIds: Set<string> | string[], message: ServerMsgT): void;
-
-  /**
-   * Send message to a specific connection.
-   */
-  sendToConnection(connectionId: string, message: ServerMsgT): void;
-}
-
-/**
- * Interface for processing messages from the client.  Also handles guest
- * admission.
- */
-export interface IMessageProcessor<ClientMsgT> {
-  authenticate(token: string): Promise<string | undefined>;
-
-  processMessage(
-    connectionId: string,
-    userId: string,
-    message: ClientMsgT
-  ): Promise<void>;
-
-  /**
-   * Handle user disconnect - clean up from all sessions.
-   * Called when a connection is closed to ensure proper cleanup.
-   */
-  handleUserDisconnect(userId: string): Promise<void>;
-}
-
-class Connection {
-  public ws: ws.WebSocket;
-  public userId: string;
-  public heartbeat: NodeJS.Timeout;
-  public lastMessageTime: number;
-
-  constructor(ws: ws.WebSocket, userId: string, heartbeat: NodeJS.Timeout) {
-    this.ws = ws;
-    this.userId = userId;
-    this.heartbeat = heartbeat;
-    this.lastMessageTime = Date.now();
-  }
-}
-
-/**
- * Manages WebSocket connections and handles protocol-level message routing.
- *
- * - Connection handshake
- * - Connection lifecycle management
- * - Message routing between connections and sessions
- * - Connection-to-user
- */
-export class ConnectionManager<ClientMsgT, ServerMsgT>
-  implements IUserConnectionManager<ServerMsgT>
-{
-  // Connections by id
-  private connections: Map<string, Connection> = new Map();
-  // userId -> connectionIds
-  private userToConnections: Map<string, Set<string>> = new Map();
-
-  // this is a singleton object and can only be initialized by .init method
-  private constructor(
-    private sessionMessageProcessor: IMessageProcessor<ClientMsgT>
-  ) {}
-
-  static init<ClientMsgT, ServerMsgT>(
-    createMsgProcessor: (
-      cm: IUserConnectionManager<ServerMsgT>
-    ) => IMessageProcessor<ClientMsgT>
-  ): ConnectionManager<ClientMsgT, ServerMsgT> {
-    const connManager = new ConnectionManager(
-      {} as IMessageProcessor<ClientMsgT>
-    );
-    connManager.sessionMessageProcessor = createMsgProcessor(connManager);
-    return connManager;
-  }
-
-  /**
-   * Handle new WebSocket connection with handshake protocol
-   * TODO: proper error code returning
-   */
-  async handleConnection(
-    webSocket: ws.WebSocket,
-    token: string,
-    req?: { headers: { [key: string]: string | string[] | undefined } }
-  ): Promise<void> {
-    const connectionId = uuidv4();
-    logger.info(`[ConnectionManager] New connection: ${connectionId}`);
-
-    try {
-      const userUUID = await this.sessionMessageProcessor.authenticate(token);
-      if (!userUUID) {
-        throw new ChatFatalError("invalid api key");
-      }
-
-      // Heartbeat interval
-      const heartbeat = setInterval(
-        this.onHeartbeat.bind(this, connectionId),
-        HEARTBEAT_INTERVAL_MS
-      );
-      // Register connection
-      this.connections.set(
-        connectionId,
-        new Connection(webSocket, userUUID, heartbeat)
-      );
-
-      // Add to user connections
-      if (!this.userToConnections.has(userUUID)) {
-        this.userToConnections.set(userUUID, new Set());
-      }
-      const userConnections = this.userToConnections.get(userUUID);
-      if (userConnections) {
-        userConnections.add(connectionId);
-      }
-
-      // Setup connection handlers
-      this.setupConnectionHandlers(connectionId, webSocket);
-
-      // Send connection_ready after WebSocket is fully open
-      const clientVersion =
-        (req?.headers["x-client-version"] as string) || "unknown";
-
-      const response: ServerConnectionReady = {
-        t: "ready",
-        c_id: connectionId,
-        user_uuid: userUUID,
-      };
-
-      // Check if connection still exists before sending
-      if (this.connections.has(connectionId)) {
-        this.sendConnectionReadyMessage(connectionId, response);
-      } else {
-        logger.warn(
-          `[ConnectionManager] Connection ${connectionId} was closed` +
-            ` before ready message could be sent`
-        );
-      }
-
-      logger.info(
-        `[ConnectionManager] Connection ${connectionId} registered ` +
-          `for user ${userUUID}, version: ${clientVersion}`
-      );
-    } catch (error) {
-      const errorMessage =
-        error instanceof Error ? error.message : String(error);
-      logger.error(
-        `[ConnectionManager] Failed to setup connection ${connectionId}:`,
-        errorMessage
-      );
-
-      // Send error to client if possible
-      this.sendErrorDirect(webSocket, connectionId, errorMessage);
-
-      // Force cleanup of any partial state
-      this.forceCleanupConnection(connectionId);
-
-      // Close WebSocket - this might trigger handlers if they were attached,
-      // but we've already cleaned up state above as a safeguard
-      try {
-        webSocket.close(4000, "Connection setup failed");
-      } catch (closeError) {
-        logger.error(
-          `[ConnectionManager] Error closing WebSocket ${connectionId}:`,
-          closeError
-        );
-      }
-    }
-  }
-
-  /**
-   * Send connection_ready message with unified error handling
-   */
-  private sendConnectionReadyMessage(
-    connectionId: string,
-    response: ServerConnectionReady
-  ): void {
-    try {
-      this.sendConnectionMessage(connectionId, response);
-    } catch (error) {
-      logger.error(
-        `[ConnectionManager] Failed connection_ready to ${connectionId}:`,
-        error
-      );
-      this.forceCleanupConnection(connectionId);
-    }
-  }
-
-  /**
-   * Send error message to connection directly.
-   * This is used when the connection is possibly not yet registered.
-   */
-  private sendErrorDirect(
-    ws: ws.WebSocket,
-    connectionId: string,
-    errorMessage: string
-  ): void {
-    const errorMsg: ServerConnectionError = {
-      t: "error",
-      e: errorMessage,
-    };
-    try {
-      if (ws.readyState === ws.OPEN) {
-        ws.send(JSON.stringify(errorMsg));
-      } else {
-        logger.error(
-          `[ConnectionManager] Conn not open sending error ${connectionId}`
-        );
-      }
-    } catch (error) {
-      logger.error(
-        `[ConnectionManager] Failed to send error to connection` +
-          ` ${connectionId}:`,
-        error
-      );
-    }
-  }
-
-  /**
-   * Setup message and lifecycle handlers for a connection
-   */
-  private setupConnectionHandlers(
-    connectionId: string,
-    webSocket: ws.WebSocket
-  ): void {
-    // TODO: keep the Connection object in the closure here instead of always
-    // doing many (expensive) lookups.
-
-    // set message handler for incoming messages
-    webSocket.on("message", (data: ws.RawData) => {
-      void (async () => {
-        try {
-          const msgStr = (data as Buffer).toString("utf8");
-          logger.info(
-            `[ConnectionManager] Message from ${connectionId}: ${msgStr}`
-          );
-          const message = JSON.parse(
-            msgStr
-          ) as ClientConnectionMessage<ClientMsgT>;
-          await this.routeMessage(connectionId, message);
-        } catch (error) {
-          logger.error(`[ConnectionManager] Error processing message:`, error);
-          this.handleError(connectionId, error);
-        }
-      })();
-    });
-
-    webSocket.on("close", () => {
-      logger.debug(`[ConnectionManager] Connection ${connectionId} closed`);
-      this.handleConnectionClose(connectionId);
-    });
-
-    webSocket.on("error", (error) => {
-      logger.error(
-        `[ConnectionManager] WebSocket error on ${connectionId}:`,
-        error
-      );
-      this.handleConnectionClose(connectionId);
-    });
-  }
-
-  /**
-   * Route incoming message to appropriate handler based on message type
-   */
-  private async routeMessage(
-    connectionId: string,
-    msg: ClientConnectionMessage<ClientMsgT>
-  ): Promise<void> {
-    const conn = this.connections.get(connectionId);
-    if (!conn) {
-      throw new ChatErrorMessage(
-        `Connection ${connectionId} not found for message routing`
-      );
-    }
-
-    conn.lastMessageTime = Date.now();
-
-    // Handle Connection-level messages
-    {
-      switch (msg.t) {
-        case "data":
-          await this.sessionMessageProcessor.processMessage(
-            connectionId,
-            conn.userId,
-            msg.d
-          );
-          break;
-        case "pong":
-          // The only operation is to update the `lastMessageTime`, done above
-          logger.debug(
-            `[ConnectionManager.routeMessage] pong from ${connectionId}`
-          );
-          break;
-        default: {
-          const _: never = msg;
-          break;
-        }
-      }
-    }
-  }
-
-  private onHeartbeat(connectionId: string): void {
-    // TODO: move this and other handling to the Connection class so we can
-    // avoid all the per-event lookups.
-
-    const conn = this.connections.get(connectionId);
-    if (!conn) {
-      // TODO: We intentionally don't try to clean up here, to force
-      // correctness fo the real cleanup code.
-      logger.error(
-        `[onHeartbeat] unexpected callback for connId ${connectionId}`
-      );
-      return;
-    }
-
-    const now = Date.now();
-    const timeSinceLast = now - conn.lastMessageTime;
-    logger.debug(
-      `[onHeartbeat] conn ${connectionId} since ${String(timeSinceLast)}`
-    );
-    if (timeSinceLast > 2 * HEARTBEAT_INTERVAL_MS) {
-      this.forceCleanupConnection(connectionId);
-      return;
-    }
-
-    const msg: ServerConnectionPing = { t: "ping" };
-    this.sendConnectionMessage(connectionId, msg);
-  }
-
-  /**
-   * Force cleanup of connection state - used for error recovery.
-   * This method is safe to call multiple times and handles partial state.
-   */
-  private forceCleanupConnection(connectionId: string): void {
-    const conn = this.connections.get(connectionId);
-    if (conn) {
-      // Cancel the heartbeat interval
-      clearInterval(conn.heartbeat);
-
-      // Remove from user connections
-      const userConnections = this.userToConnections.get(conn.userId);
-      if (userConnections) {
-        userConnections.delete(connectionId);
-        if (userConnections.size === 0) {
-          this.userToConnections.delete(conn.userId);
-          const _ = this.sessionMessageProcessor.handleUserDisconnect(
-            conn.userId
-          );
-        }
-      }
-
-      this.connections.delete(connectionId);
-    }
-
-    logger.debug(
-      `[ConnectionManager] Force cleaned up connection ${connectionId}`
-    );
-  }
-
-  /**
-   * Handle connection close - cleanup all state.
-   * This is the normal cleanup path triggered by WebSocket events.
-   */
-  private handleConnectionClose(connectionId: string): void {
-    this.forceCleanupConnection(connectionId);
-    logger.info(`[ConnectionManager] Connection ${connectionId} closed`);
-  }
-
-  private sendConnectionMessage(
-    connectionId: string,
-    message: ServerConnectionMessage<ServerMsgT>
-  ): void {
-    const conn = this.connections.get(connectionId);
-    if (!conn) {
-      logger.warn(
-        `[ConnectionManager] Cannot send to connection ${connectionId} - ` +
-          `not found`
-      );
-      return;
-    }
-
-    const webSocket = conn.ws;
-    if (webSocket.readyState !== 1) {
-      // WebSocket.OPEN = 1
-      logger.warn(
-        `[ConnectionManager] Cannot send to connection ${connectionId} - ` +
-          `not open`
-      );
-      return;
-    }
-
-    const msgString = JSON.stringify(message);
-    logger.debug(
-      `[ConnectionManager] Sending to ${connectionId}: ${msgString}`
-    );
-    webSocket.send(msgString);
-  }
-
-  /**
-   * Send message to specific connection.
-   */
-  public sendToConnection(connectionId: string, message: ServerMsgT): void {
-    this.sendConnectionMessage(connectionId, { t: "data", d: message });
-  }
-
-  /**
-   * Send error message to connection
-   */
-  private sendConnectionError(
-    connectionId: string,
-    errorMessage: string
-  ): void {
-    const errorMsg: ServerConnectionError = {
-      t: "error",
-      e: errorMessage,
-    };
-    this.sendConnectionMessage(connectionId, errorMsg);
-  }
-
-  /**
-   * Handle errors during message processing
-   */
-  private handleError(connectionId: string, error: unknown): void {
-    let message: string;
-    let shouldClose = false;
-
-    if (typeof error === "string") {
-      message = error;
-    } else if (error instanceof ChatFatalError) {
-      message = error.message;
-      shouldClose = true;
-    } else if (error instanceof ChatErrorMessage) {
-      message = error.message;
-    } else if (error instanceof Error) {
-      message = "Internal server error: " + error.message;
-    } else {
-      message = "Unknown error occurred";
-    }
-
-    logger.warn(
-      `[ConnectionManager] Error on connection ${connectionId}: ${message}`
-    );
-    this.sendConnectionError(connectionId, message);
-
-    if (shouldClose) {
-      const conn = this.connections.get(connectionId);
-      logger.info(
-        `[ConnectionManager] Closing connection ${connectionId} ` +
-          `due to error: ${message}`
-      );
-      if (conn) {
-        conn.ws.close(4000, message);
-      }
-    }
-  }
-
-  /**
-   * Implementation of IUserConnectionManager interface.
-   * Send message to all active connections of specific users.
-   */
-  public sendToUsers(userIds: Set<string>, message: ServerMsgT): void {
-    for (const userId of userIds.values()) {
-      const connectionIds = this.getUserConnections(userId);
-      for (const connectionId of connectionIds) {
-        logger.debug(
-          `[ConnectionManager] sending to user ${userId} ` +
-            `connection ${connectionId}`
-        );
-        this.sendToConnection(connectionId, message);
-      }
-    }
-  }
-
-  /**
-   * Get all active connection IDs for a user.
-   * Helper method for sendToUsers implementation.
-   */
-  private getUserConnections(userId: string): string[] {
-    const connections = this.userToConnections.get(userId);
-    return connections ? Array.from(connections) : [];
-  }
-}