@xalia/agent 0.5.8 → 0.6.0

package/dist/agent/src/chat/server/connectionManager.js

@@ -0,0 +1,352 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.ConnectionManager = void 0;
+const uuid_1 = require("uuid");
+const sdk_1 = require("@xalia/xmcp/sdk");
+const errors_1 = require("../protocol/errors");
+const logger = (0, sdk_1.getLogger)();
+const HEARTBEAT_INTERVAL_MS = parseInt(process.env["HEARTBEAT_INTERVAL_MS"] || String(30 * 1000));
+class Connection {
+    constructor(ws, userId, apiKey, heartbeat) {
+        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
+ */
+class ConnectionManager {
+    // this is a singleton object and can only be initialized by .init method
+    constructor(apiKeyManager, sessionMessageProcessor) {
+        this.apiKeyManager = apiKeyManager;
+        this.sessionMessageProcessor = sessionMessageProcessor;
+        // Connections by id
+        this.connections = new Map();
+        // userId -> connectionIds
+        this.userToConnections = new Map();
+    }
+    static init(apiKeyManager, createMsgProcessor) {
+        const connManager = new ConnectionManager(apiKeyManager, {});
+        connManager.sessionMessageProcessor = createMsgProcessor(connManager);
+        return connManager;
+    }
+    /**
+     * Handle new WebSocket connection with handshake protocol
+     * TODO: proper error code returning
+     */
+    async handleConnection(webSocket, apiKey, req) {
+        const connectionId = (0, uuid_1.v4)();
+        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 errors_1.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"] || "unknown";
+            const response = {
+                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
+     */
+    sendConnectionReadyMessage(connectionId, response) {
+        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.
+     */
+    sendErrorDirect(ws, connectionId, errorMessage) {
+        const errorMsg = {
+            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
+     */
+    setupConnectionHandlers(connectionId, webSocket) {
+        // 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) => {
+            void (async () => {
+                try {
+                    const msgStr = data.toString("utf8");
+                    logger.info(`[ConnectionManager] Message from ${connectionId}: ${msgStr}`);
+                    const message = JSON.parse(msgStr);
+                    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
+     */
+    async routeMessage(connectionId, msg) {
+        const conn = this.connections.get(connectionId);
+        if (!conn) {
+            throw new errors_1.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 _ = msg;
+                    break;
+                }
+            }
+        }
+    }
+    onHeartbeat(connectionId) {
+        // 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 = { 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.
+     */
+    forceCleanupConnection(connectionId) {
+        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.
+     */
+    handleConnectionClose(connectionId) {
+        this.forceCleanupConnection(connectionId);
+        logger.info(`[ConnectionManager] Connection ${connectionId} closed`);
+    }
+    sendConnectionMessage(connectionId, message) {
+        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.
+     */
+    sendToConnection(connectionId, message) {
+        this.sendConnectionMessage(connectionId, { t: "data", d: message });
+    }
+    /**
+     * Send error message to connection
+     */
+    sendConnectionError(connectionId, errorMessage) {
+        const errorMsg = {
+            t: "error",
+            e: errorMessage,
+        };
+        this.sendConnectionMessage(connectionId, errorMsg);
+    }
+    /**
+     * Handle errors during message processing
+     */
+    handleError(connectionId, error) {
+        let message;
+        let shouldClose = false;
+        if (typeof error === "string") {
+            message = error;
+        }
+        else if (error instanceof errors_1.ChatFatalError) {
+            message = error.message;
+            shouldClose = true;
+        }
+        else if (error instanceof errors_1.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.
+     */
+    sendToUsers(userIds, message) {
+        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.
+     */
+    getUserConnections(userId) {
+        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.
+     */
+    getLiveUserApiKey(userId) {
+        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;
+    }
+}
+exports.ConnectionManager = ConnectionManager;

package/dist/agent/src/chat/server/connectionManager.test.js

@@ -0,0 +1,159 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+/* eslint-disable @typescript-eslint/no-unsafe-assignment */
+/* eslint-disable @typescript-eslint/no-unsafe-call */
+const vitest_1 = require("vitest");
+const connectionManager_1 = require("./connectionManager");
+const mockFactories_1 = require("./test-utils/mockFactories");
+// Mock dependencies
+vitest_1.vi.mock("../data/apiKeyManager");
+vitest_1.vi.mock("ws");
+(0, vitest_1.describe)("ConnectionManager", () => {
+    let connectionManager;
+    let mockApiKeyManager;
+    let mockSessionRegistry;
+    let mockWebSocket;
+    (0, vitest_1.beforeEach)(() => {
+        vitest_1.vi.resetAllMocks();
+        // Create mocks using factories
+        mockApiKeyManager = (0, mockFactories_1.createMockApiKeyManager)();
+        mockSessionRegistry = (0, mockFactories_1.createMockSessionRegistry)();
+        mockWebSocket = (0, mockFactories_1.createMockWebSocket)();
+        // Setup standard behaviors
+        (0, mockFactories_1.setupStandardMockBehaviors)({
+            apiKeyManager: mockApiKeyManager,
+        });
+        // Create ConnectionManager instance using reflection
+        // to access private constructor. Since the constructor is private,
+        // we need to use type assertion
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        connectionManager = new connectionManager_1.ConnectionManager(mockApiKeyManager.mock, mockSessionRegistry.mock);
+    });
+    (0, vitest_1.afterEach)(() => {
+        vitest_1.vi.restoreAllMocks();
+    });
+    (0, vitest_1.describe)("handleConnection", () => {
+        (0, vitest_1.it)("should send connection_ready after successful setup", async () => {
+            await connectionManager.handleConnection(mockWebSocket.mock, "valid-api-key", { headers: {} });
+            // Wait for deferred send
+            await new Promise((resolve) => setImmediate(resolve));
+            (0, vitest_1.expect)(mockApiKeyManager.spies.verifyApiKey).toHaveBeenCalledWith("valid-api-key");
+            (0, vitest_1.expect)(mockWebSocket.spies.send).toHaveBeenCalledWith(vitest_1.expect.stringMatching(/"t":"ready"/));
+        });
+        (0, vitest_1.it)("should reject connection with invalid API key", async () => {
+            await connectionManager.handleConnection(mockWebSocket.mock, "invalid-api-key");
+            (0, vitest_1.expect)(mockWebSocket.spies.close).toHaveBeenCalledWith(4000, "Connection setup failed");
+        });
+        (0, vitest_1.it)("should setup WebSocket event handlers", async () => {
+            await connectionManager.handleConnection(mockWebSocket.mock, "valid-api-key");
+            (0, vitest_1.expect)(mockWebSocket.spies.on).toHaveBeenCalledWith("message", vitest_1.expect.any(Function));
+            (0, vitest_1.expect)(mockWebSocket.spies.on).toHaveBeenCalledWith("close", vitest_1.expect.any(Function));
+            (0, vitest_1.expect)(mockWebSocket.spies.on).toHaveBeenCalledWith("error", vitest_1.expect.any(Function));
+        });
+    });
+    (0, vitest_1.describe)("user communication", () => {
+        (0, vitest_1.beforeEach)(async () => {
+            await connectionManager.handleConnection(mockWebSocket.mock, "valid-api-key");
+        });
+        (0, vitest_1.it)("should send messages to connected users", () => {
+            const testMessage = {
+                t: "data",
+                d: {
+                    type: "user_msg",
+                    user_uuid: "user_1",
+                    session_id: "session_1",
+                    message: "message text",
+                    message_idx: 12,
+                    // connection_id: "test-123",
+                    // user_uuid: MOCK_USERS.owner.uuid,
+                    // client_message_id: "test-123",
+                },
+            };
+            connectionManager.sendToUsers(new Set([mockFactories_1.MOCK_USERS.owner.uuid]), testMessage.d);
+            (0, vitest_1.expect)(mockWebSocket.spies.send).toHaveBeenCalledWith(JSON.stringify(testMessage));
+        });
+        (0, vitest_1.it)("should return API key for connected user", () => {
+            const apiKey = connectionManager.getLiveUserApiKey(mockFactories_1.MOCK_USERS.owner.uuid);
+            (0, vitest_1.expect)(apiKey).toBe("valid-api-key");
+        });
+        (0, vitest_1.it)("should return undefined for non-existent user", () => {
+            const apiKey = connectionManager.getLiveUserApiKey(mockFactories_1.MOCK_USERS.nonExistent.uuid);
+            (0, vitest_1.expect)(apiKey).toBeUndefined();
+        });
+    });
+    (0, vitest_1.describe)("connection cleanup", () => {
+        (0, vitest_1.beforeEach)(async () => {
+            await connectionManager.handleConnection(mockWebSocket.mock, "valid-api-key");
+        });
+        (0, vitest_1.it)("should clean up connection state on close", () => {
+            // Get the close handler from mock calls
+            const closeHandlerCall = mockWebSocket.spies.on.mock.calls.find((call) => call[0] === "close");
+            (0, vitest_1.expect)(closeHandlerCall).toBeDefined();
+            if (closeHandlerCall) {
+                const closeHandler = closeHandlerCall[1];
+                closeHandler();
+            }
+            // After cleanup, API key should no longer be available
+            const apiKey = connectionManager.getLiveUserApiKey(mockFactories_1.MOCK_USERS.owner.uuid);
+            (0, vitest_1.expect)(apiKey).toBeUndefined();
+        });
+        (0, vitest_1.it)("should clean up connection state on error", () => {
+            // Get the error handler from mock calls
+            const errorHandlerCall = mockWebSocket.spies.on.mock.calls.find((call) => call[0] === "error");
+            (0, vitest_1.expect)(errorHandlerCall).toBeDefined();
+            if (errorHandlerCall) {
+                const errorHandler = errorHandlerCall[1];
+                errorHandler(new Error("WebSocket error"));
+            }
+            // After cleanup, API key should no longer be available
+            const apiKey = connectionManager.getLiveUserApiKey(mockFactories_1.MOCK_USERS.owner.uuid);
+            (0, vitest_1.expect)(apiKey).toBeUndefined();
+        });
+    });
+    (0, vitest_1.describe)("message routing", () => {
+        let messageHandler;
+        (0, vitest_1.beforeEach)(async () => {
+            await connectionManager.handleConnection(mockWebSocket.mock, "valid-api-key");
+            // Extract the message handler
+            const messageHandlerCall = mockWebSocket.spies.on.mock.calls.find((call) => call[0] === "message");
+            (0, vitest_1.expect)(messageHandlerCall).toBeDefined();
+            if (!messageHandlerCall) {
+                throw new Error("Message handler not found");
+            }
+            messageHandler = messageHandlerCall[1];
+        });
+        (0, vitest_1.it)("should route session_create to SessionRegistry", async () => {
+            mockSessionRegistry.spies.processMessage.mockResolvedValue(undefined);
+            const message = {
+                t: "data",
+                d: {
+                    type: "control_session_create",
+                    title: "Test Session",
+                    agent_profile_id: "agent-123",
+                    client_message_id: "msg-123",
+                },
+            };
+            await messageHandler(Buffer.from(JSON.stringify(message)));
+            (0, vitest_1.expect)(mockSessionRegistry.spies.processMessage).toHaveBeenCalledWith(vitest_1.expect.any(String), // connectionId
+            mockFactories_1.MOCK_USERS.owner.uuid, message.d);
+        });
+        (0, vitest_1.it)("should route session messages to SessionRegistry", async () => {
+            const message = {
+                t: "data",
+                d: {
+                    type: "msg",
+                    session_id: "session-123",
+                    message: "Hello world!",
+                    client_message_id: "msg-123",
+                },
+            };
+            await messageHandler(Buffer.from(JSON.stringify(message)));
+            (0, vitest_1.expect)(mockSessionRegistry.spies.processMessage).toHaveBeenCalledWith(vitest_1.expect.any(String), // connectionId
+            mockFactories_1.MOCK_USERS.owner.uuid, message.d);
+        });
+        (0, vitest_1.it)("should handle invalid JSON gracefully", async () => {
+            await messageHandler(Buffer.from("invalid json"));
+            (0, vitest_1.expect)(mockWebSocket.spies.send).toHaveBeenCalledWith(vitest_1.expect.stringMatching(/"t":"error"/));
+        });
+    });
+});