@xalia/agent 0.5.7 → 0.6.0

package/src/chat/server/connectionManager.test.ts

@@ -0,0 +1,266 @@
+/* eslint-disable @typescript-eslint/no-unsafe-assignment */
+/* eslint-disable @typescript-eslint/no-unsafe-call */
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import { ConnectionManager } from "./connectionManager";
+import type { ClientToServer, ServerToClient } from "../protocol/messages";
+import {
+  createMockApiKeyManager,
+  createMockSessionRegistry,
+  createMockWebSocket,
+  setupStandardMockBehaviors,
+  MOCK_USERS,
+} from "./test-utils/mockFactories";
+import {
+  ClientConnectionData,
+  ServerConnectionData,
+} from "../protocol/connectionMessages";
+
+// Mock dependencies
+vi.mock("../data/apiKeyManager");
+vi.mock("ws");
+
+describe("ConnectionManager", () => {
+  let connectionManager: ConnectionManager<ClientToServer, ServerToClient>;
+  let mockApiKeyManager: ReturnType<typeof createMockApiKeyManager>;
+  let mockSessionRegistry: ReturnType<typeof createMockSessionRegistry>;
+  let mockWebSocket: ReturnType<typeof createMockWebSocket>;
+
+  beforeEach(() => {
+    vi.resetAllMocks();
+
+    // Create mocks using factories
+    mockApiKeyManager = createMockApiKeyManager();
+    mockSessionRegistry = createMockSessionRegistry();
+    mockWebSocket = createMockWebSocket();
+
+    // Setup standard behaviors
+    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 as any)(
+      mockApiKeyManager.mock,
+      mockSessionRegistry.mock
+    );
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  describe("handleConnection", () => {
+    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));
+
+      expect(mockApiKeyManager.spies.verifyApiKey).toHaveBeenCalledWith(
+        "valid-api-key"
+      );
+      expect(mockWebSocket.spies.send).toHaveBeenCalledWith(
+        expect.stringMatching(/"t":"ready"/)
+      );
+    });
+
+    it("should reject connection with invalid API key", async () => {
+      await connectionManager.handleConnection(
+        mockWebSocket.mock,
+        "invalid-api-key"
+      );
+
+      expect(mockWebSocket.spies.close).toHaveBeenCalledWith(
+        4000,
+        "Connection setup failed"
+      );
+    });
+
+    it("should setup WebSocket event handlers", async () => {
+      await connectionManager.handleConnection(
+        mockWebSocket.mock,
+        "valid-api-key"
+      );
+
+      expect(mockWebSocket.spies.on).toHaveBeenCalledWith(
+        "message",
+        expect.any(Function)
+      );
+      expect(mockWebSocket.spies.on).toHaveBeenCalledWith(
+        "close",
+        expect.any(Function)
+      );
+      expect(mockWebSocket.spies.on).toHaveBeenCalledWith(
+        "error",
+        expect.any(Function)
+      );
+    });
+  });
+
+  describe("user communication", () => {
+    beforeEach(async () => {
+      await connectionManager.handleConnection(
+        mockWebSocket.mock,
+        "valid-api-key"
+      );
+    });
+
+    it("should send messages to connected users", () => {
+      const testMessage: ServerConnectionData<ServerToClient> = {
+        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([MOCK_USERS.owner.uuid]),
+        testMessage.d
+      );
+
+      expect(mockWebSocket.spies.send).toHaveBeenCalledWith(
+        JSON.stringify(testMessage)
+      );
+    });
+
+    it("should return API key for connected user", () => {
+      const apiKey = connectionManager.getLiveUserApiKey(MOCK_USERS.owner.uuid);
+      expect(apiKey).toBe("valid-api-key");
+    });
+
+    it("should return undefined for non-existent user", () => {
+      const apiKey = connectionManager.getLiveUserApiKey(
+        MOCK_USERS.nonExistent.uuid
+      );
+      expect(apiKey).toBeUndefined();
+    });
+  });
+
+  describe("connection cleanup", () => {
+    beforeEach(async () => {
+      await connectionManager.handleConnection(
+        mockWebSocket.mock,
+        "valid-api-key"
+      );
+    });
+
+    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"
+      );
+      expect(closeHandlerCall).toBeDefined();
+
+      if (closeHandlerCall) {
+        const closeHandler = closeHandlerCall[1];
+        closeHandler();
+      }
+
+      // After cleanup, API key should no longer be available
+      const apiKey = connectionManager.getLiveUserApiKey(MOCK_USERS.owner.uuid);
+      expect(apiKey).toBeUndefined();
+    });
+
+    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"
+      );
+      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(MOCK_USERS.owner.uuid);
+      expect(apiKey).toBeUndefined();
+    });
+  });
+
+  describe("message routing", () => {
+    let messageHandler: (data: Buffer) => Promise<void>;
+
+    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"
+      );
+      expect(messageHandlerCall).toBeDefined();
+      if (!messageHandlerCall) {
+        throw new Error("Message handler not found");
+      }
+      messageHandler = messageHandlerCall[1];
+    });
+
+    it("should route session_create to SessionRegistry", async () => {
+      mockSessionRegistry.spies.processMessage.mockResolvedValue(undefined);
+
+      const message: ClientConnectionData<ClientToServer> = {
+        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)));
+
+      expect(mockSessionRegistry.spies.processMessage).toHaveBeenCalledWith(
+        expect.any(String), // connectionId
+        MOCK_USERS.owner.uuid,
+        message.d
+      );
+    });
+
+    it("should route session messages to SessionRegistry", async () => {
+      const message: ClientConnectionData<ClientToServer> = {
+        t: "data",
+        d: {
+          type: "msg",
+          session_id: "session-123",
+          message: "Hello world!",
+          client_message_id: "msg-123",
+        },
+      };
+
+      await messageHandler(Buffer.from(JSON.stringify(message)));
+
+      expect(mockSessionRegistry.spies.processMessage).toHaveBeenCalledWith(
+        expect.any(String), // connectionId
+        MOCK_USERS.owner.uuid,
+        message.d
+      );
+    });
+
+    it("should handle invalid JSON gracefully", async () => {
+      await messageHandler(Buffer.from("invalid json"));
+
+      expect(mockWebSocket.spies.send).toHaveBeenCalledWith(
+        expect.stringMatching(/"t":"error"/)
+      );
+    });
+  });
+});