@xalia/agent 0.5.8 → 0.6.0

package/src/chat/server/connectionManager.ts

@@ -0,0 +1,541 @@
+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";
+import { ApiKeyManager } from "../data/apiKeyManager";
+
+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>, message: ServerMsgT): void;
+
+  /**
+   * Send message to a specific connection.
+   */
+  sendToConnection(connectionId: string, message: ServerMsgT): void;
+
+  /**
+   * Get cached api key for a user. This key will not be available
+   * after the user disconnects.
+   */
+  getLiveUserApiKey(userId: string): string | undefined;
+}
+
+/**
+ * Interface for processing messages from the client.
+ */
+export interface IMessageProcessor<ClientMsgT> {
+  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): void;
+}
+
+class Connection {
+  public ws: ws.WebSocket;
+  public userId: string;
+  public apiKey: string;
+  public heartbeat: NodeJS.Timeout;
+  public lastMessageTime: number;
+
+  constructor(
+    ws: ws.WebSocket,
+    userId: string,
+    apiKey: string,
+    heartbeat: NodeJS.Timeout
+  ) {
+    this.ws = ws;
+    this.userId = userId;
+    this.apiKey = apiKey;
+    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 apiKeyManager: ApiKeyManager,
+    private sessionMessageProcessor: IMessageProcessor<ClientMsgT>
+  ) {}
+
+  static init<ClientMsgT, ServerMsgT>(
+    apiKeyManager: ApiKeyManager,
+    createMsgProcessor: (
+      cm: IUserConnectionManager<ServerMsgT>
+    ) => IMessageProcessor<ClientMsgT>
+  ): ConnectionManager<ClientMsgT, ServerMsgT> {
+    const connManager = new ConnectionManager(
+      apiKeyManager,
+      {} 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,
+    apiKey: string,
+    req?: { headers: { [key: string]: string | string[] | undefined } }
+  ): Promise<void> {
+    const connectionId = uuidv4();
+    logger.info(`[ConnectionManager] New connection: ${connectionId}`);
+
+    try {
+      // Verify API key
+      const userData = await this.apiKeyManager.verifyApiKey(apiKey);
+      if (!userData) {
+        logger.error(`[ConnectionManager] Invalid API key: ${apiKey}`);
+        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, userData.uuid, apiKey, heartbeat)
+      );
+
+      // Add to user connections
+      if (!this.userToConnections.has(userData.uuid)) {
+        this.userToConnections.set(userData.uuid, new Set());
+      }
+      const userConnections = this.userToConnections.get(userData.uuid);
+      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: userData.uuid,
+      };
+
+      // 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 ${userData.uuid}, 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);
+          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) : [];
+  }
+
+  /**
+   * Implementation of IUserConnectionManager interface.
+   * Get cached api key for a user. This key will not be available
+   * after the user disconnects.
+   */
+  public getLiveUserApiKey(userId: string): string | undefined {
+    const connectionIds = this.getUserConnections(userId);
+
+    // Return API key from the first active connection
+    // (all connections for a user should have the same API key)
+    for (const connectionId of connectionIds) {
+      const conn = this.connections.get(connectionId);
+      if (conn) {
+        return conn.apiKey;
+      }
+    }
+
+    // No active connections for this user
+    return undefined;
+  }
+}