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

import { describe, it, expect, vi } from "vitest";

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

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?.();
  }
  receive(frame: unknown): void {
    this.onmessage?.({ data: JSON.stringify(frame) });
  }
}

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 };
}

function attachedAck(id: string, sessionId: string) {
  return {
    kind: "terminal.attached",
    replyTo: id,
    payload: {
      session: { sessionId, rows: 30, cols: 120 },
      scrollback: [],
      nextSeq: 0,
      status: "running",
      gap: false,
    },
  };
}

async function attachTerminal(
  client: WsLiveClient,
  sockets: FakeSocket[],
  sessionId: string,
): Promise<void> {
  const attached = client.attachTerminal({ sessionId, onData: () => {} });
  await vi.waitFor(() => expect(sockets[0]?.sent.length ?? 0).toBeGreaterThan(0));
  const sent = JSON.parse(sockets[0].sent[0]);
  sockets[0].receive(attachedAck(sent.id, sessionId));
  await attached;
}

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

    const received: unknown[] = [];
    client.setChatSink("chat-1", (chunk) => received.push(chunk));

    socket.receive({
      kind: "chat.output",
      payload: {
        sessionId: "chat-1",
        chunk: { kind: "final", content: "done" },
      },
    });

    expect(received).toEqual([{ kind: "final", content: "done" }]);

    client.removeChatSink("chat-1");
    socket.receive({
      kind: "chat.output",
      payload: {
        sessionId: "chat-1",
        chunk: { kind: "final", content: "ignored" },
      },
    });
    expect(received).toEqual([{ kind: "final", content: "done" }]);
  });

  it("clears structured chat sinks on disconnect", async () => {
    const { client, socket } = connectedClient();
    await client.ensureConnected();

    const received: unknown[] = [];
    client.setChatSink("chat-1", (chunk) => received.push(chunk));
    client.disconnect();
    expect(client.getConnectionState()).toBe("closed");

    socket.receive({
      kind: "chat.output",
      payload: {
        sessionId: "chat-1",
        chunk: { kind: "final", content: "dropped" },
      },
    });
    expect(received).toEqual([]);
  });

  it("re-attaches chat sinks after reconnect while PTY logs keep the socket alive", async () => {
    const sockets: FakeSocket[] = [];
    let timerFn: (() => void) | null = null;
    const client = new WsLiveClient({
      wsUrl: "wss://host",
      reconnectBaseMs: 1,
      socketFactory: () => {
        const socket = new FakeSocket();
        sockets.push(socket);
        queueMicrotask(() => socket.onopen?.());
        return socket;
      },
      setTimeoutImpl: (fn) => {
        timerFn = fn;
        return 1;
      },
      clearTimeoutImpl: () => {
        timerFn = null;
      },
    });

    await attachTerminal(client, sockets, "pty-1");
    const received: unknown[] = [];
    client.setChatSink("chat-1", (chunk) => received.push(chunk));

    sockets[0].close();
    expect(client.getConnectionState()).toBe("reconnecting");
    expect(client.needsReconnect()).toBe(true);

    const fire = timerFn;
    timerFn = null;
    fire?.();
    await vi.waitFor(() => expect(sockets.length).toBe(2));
    await vi.waitFor(() => expect(sockets[1].sent.length).toBeGreaterThan(0));

    const terminalAttach = JSON.parse(sockets[1].sent[0]);
    expect(terminalAttach.kind).toBe("terminal.attach");
    sockets[1].receive(attachedAck(terminalAttach.id, "pty-1"));

    await vi.waitFor(() => expect(sockets[1].sent.length).toBeGreaterThan(1));
    const chatAttach = JSON.parse(sockets[1].sent[1]);
    expect(chatAttach.kind).toBe("chat.attach");
    expect(chatAttach.payload).toEqual({ sessionId: "chat-1" });
    sockets[1].receive({
      kind: "chat.attached",
      replyTo: chatAttach.id,
      payload: { sessionId: "chat-1", cellKind: "chat" },
    });

    sockets[1].receive({
      kind: "terminal.output",
      payload: {
        sessionId: "pty-1",
        seq: 1,
        bytesBase64: bytesToBase64(new Uint8Array([108, 111, 103])),
      },
    });
    sockets[1].receive({
      kind: "chat.output",
      payload: {
        sessionId: "chat-1",
        chunk: {
          kind: "progress",
          progress: {
            source: "ideaLocal",
            kind: "mcp",
            stage: "started",
            label: "idea_ticket_read",
            text: "reading ticket #241",
          },
        },
      },
    });

    expect(received).toEqual([
      {
        kind: "progress",
        progress: {
          source: "ideaLocal",
          kind: "mcp",
          stage: "started",
          label: "idea_ticket_read",
          text: "reading ticket #241",
        },
      },
    ]);
  });
});