IdeA/frontend/src/adapters/http/wsLiveClient.test.ts

/**
 * F1 — the WS live client skeleton: connects lazily, correlates a command with
 * its `replyTo` acknowledgement, routes per-session PTY output to the registered
 * sink, and dispatches `event.domain` frames to the domain-event handler.
 */
import { describe, it, expect, vi } from "vitest";

import type { WebSocketLike } from "./wsLiveClient";
import { WsLiveClient } from "./wsLiveClient";
import { bytesToBase64, base64ToBytes } from "./frames";

/** A controllable fake WebSocket that auto-completes its open handshake. */
class FakeSocket implements WebSocketLike {
  sent: string[] = [];
  onopen: (() => void) | null = null;
  onmessage: ((event: { data: string }) => void) | null = null;
  onerror: ((event: unknown) => void) | null = null;
  onclose: (() => void) | null = null;

  send(data: string): void {
    this.sent.push(data);
  }
  close(): void {
    this.onclose?.();
  }
  /** Simulate the server pushing a frame. */
  receive(frame: unknown): void {
    this.onmessage?.({ data: JSON.stringify(frame) });
  }
}

/**
 * Builds a client whose socket auto-opens on the next microtask (after the
 * client has assigned `onopen`), so `ensureConnected()`/`send()` resolve without
 * manual driving.
 */
function connectedClient(): { client: WsLiveClient; socket: FakeSocket } {
  const socket = new FakeSocket();
  const client = new WsLiveClient({
    wsUrl: "wss://host",
    socketFactory: () => {
      queueMicrotask(() => socket.onopen?.());
      return socket;
    },
  });
  return { client, socket };
}

describe("base64 helpers round-trip binary", () => {
  it("encodes and decodes arbitrary bytes", () => {
    const bytes = new Uint8Array([0, 13, 27, 255, 128, 10]);
    expect(base64ToBytes(bytesToBase64(bytes))).toEqual(bytes);
  });
});

describe("WsLiveClient", () => {
  it("connects lazily and resolves a command with its replyTo ack", async () => {
    const { client, socket } = connectedClient();

    const ackPromise = client.send("terminal.open", { cwd: "/root", rows: 30, cols: 120 });
    await vi.waitFor(() => expect(socket.sent).toHaveLength(1));

    const sent = JSON.parse(socket.sent[0]);
    expect(sent.kind).toBe("terminal.open");
    expect(sent.payload).toMatchObject({ cwd: "/root", rows: 30, cols: 120 });

    socket.receive({
      kind: "terminal.attached",
      replyTo: sent.id,
      payload: { session: { sessionId: "s1", rows: 30, cols: 120 }, nextSeq: 1, scrollback: [], gap: false },
    });

    const ack = await ackPromise;
    expect(ack.kind).toBe("terminal.attached");
  });

  it("routes terminal.output to the per-session sink", async () => {
    const { client, socket } = connectedClient();
    await client.ensureConnected();

    const received: Uint8Array[] = [];
    client.setOutputSink("s1", (bytes) => received.push(bytes));

    socket.receive({
      kind: "terminal.output",
      payload: { sessionId: "s1", seq: 1, bytesBase64: bytesToBase64(new Uint8Array([104, 105])) },
    });

    expect(received).toHaveLength(1);
    expect(received[0]).toEqual(new Uint8Array([104, 105]));

    // After removing the sink, further output is dropped.
    client.removeOutputSink("s1");
    socket.receive({
      kind: "terminal.output",
      payload: { sessionId: "s1", seq: 2, bytesBase64: bytesToBase64(new Uint8Array([106])) },
    });
    expect(received).toHaveLength(1);
  });

  it("dispatches event.domain frames to the domain-event handler", async () => {
    const { client, socket } = connectedClient();
    await client.ensureConnected();

    const handler = vi.fn();
    client.setDomainEventHandler(handler);

    socket.receive({ kind: "event.domain", payload: { type: "agentBusyChanged", agentId: "a1", busy: true } });

    expect(handler).toHaveBeenCalledWith({ type: "agentBusyChanged", agentId: "a1", busy: true });
  });

  it("rejects a pending command when an error frame arrives", async () => {
    const { client, socket } = connectedClient();
    const promise = client.send("terminal.close", { sessionId: "gone" });
    await vi.waitFor(() => expect(socket.sent).toHaveLength(1));

    const sent = JSON.parse(socket.sent[0]);
    socket.receive({ kind: "error", replyTo: sent.id, payload: { code: "NOT_FOUND", message: "no session" } });

    await expect(promise).rejects.toEqual({ code: "NOT_FOUND", message: "no session" });
  });
});