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

/**
 * F3 — the WS live client connection state machine + reconnection:
 *  - disconnected (socket close) ⇒ `reconnecting`, a notice is written to xterm;
 *  - on reconnect, the tracked terminal is re-attached with its `lastSeq` and the
 *    bounded scrollback is repainted; state returns to `connected`;
 *  - a `terminal.status` `exited` frame notifies the session and stops tracking
 *    (so a later disconnect does not try to re-attach a dead PTY).
 */
import { describe, it, expect, vi } from "vitest";

import { WsLiveClient, type WebSocketLike, type ConnectionState } 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 harness() {
  const sockets: FakeSocket[] = [];
  let timerFn: (() => void) | null = null;
  const ws = new WsLiveClient({
    wsUrl: "wss://host",
    reconnectBaseMs: 1,
    socketFactory: () => {
      const s = new FakeSocket();
      sockets.push(s);
      queueMicrotask(() => s.onopen?.());
      return s;
    },
    setTimeoutImpl: (fn) => {
      timerFn = fn;
      return 1;
    },
    clearTimeoutImpl: () => {
      timerFn = null;
    },
  });
  const states: ConnectionState[] = [];
  ws.onConnectionStateChange((s) => states.push(s));
  return {
    ws,
    sockets,
    states,
    fireTimer: () => {
      const fn = timerFn;
      timerFn = null;
      fn?.();
    },
    hasTimer: () => timerFn !== null,
  };
}

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

async function attach(
  ws: WsLiveClient,
  sockets: FakeSocket[],
  sessionId: string,
  onData: (b: Uint8Array) => void,
  onStatus?: (s: string, c?: number | null) => void,
): Promise<void> {
  const p = ws.attachTerminal({ sessionId, onData, onStatus });
  // The socket is created lazily by the factory inside ensureConnected.
  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, 0));
  await p;
}

const decode = (b: Uint8Array) => new TextDecoder().decode(b);

describe("WsLiveClient connection state + reconnection", () => {
  it("goes connecting → connected on the first attach", async () => {
    const h = harness();
    await attach(h.ws, h.sockets, "s1", () => {});
    expect(h.ws.getConnectionState()).toBe("connected");
    expect(h.states).toContain("connected");
  });

  it("on socket close it reconnects, re-attaches with lastSeq and repaints scrollback", async () => {
    const h = harness();
    const chunks: Uint8Array[] = [];
    await attach(h.ws, h.sockets, "s1", (b) => chunks.push(b));

    // Advance lastSeq via an output frame (seq 5).
    h.sockets[0].receive({
      kind: "terminal.output",
      payload: { sessionId: "s1", seq: 5, bytesBase64: bytesToBase64(new Uint8Array([120])) },
    });

    // Drop the socket → reconnecting + a "déconnecté" notice.
    h.sockets[0].close();
    expect(h.ws.getConnectionState()).toBe("reconnecting");
    expect(chunks.some((c) => decode(c).includes("déconnecté"))).toBe(true);
    expect(h.hasTimer()).toBe(true);

    // Fire the reconnect timer → a new socket opens and re-attaches.
    h.fireTimer();
    await vi.waitFor(() => expect(h.sockets.length).toBe(2));
    await vi.waitFor(() => expect(h.sockets[1].sent.length).toBeGreaterThan(0));
    const reattach = JSON.parse(h.sockets[1].sent[0]);
    expect(reattach.kind).toBe("terminal.attach");
    // lastSeq carries the last output seq seen (5) for precise replay.
    expect(reattach.payload).toMatchObject({ sessionId: "s1", lastSeq: 5 });

    // Server replies with fresh scrollback; the client repaints it + "reconnecté".
    const scroll = new Uint8Array([82, 69]);
    h.sockets[1].receive(attachedAck(reattach.id, "s1", 3, scroll));
    await vi.waitFor(() => expect(h.ws.getConnectionState()).toBe("connected"));
    // The repaint happens after the re-attach ack resolves (a microtask later).
    await vi.waitFor(() =>
      expect(chunks.some((c) => decode(c).includes("reconnecté"))).toBe(true),
    );
    expect(chunks.some((c) => c.length === 2 && c[0] === 82 && c[1] === 69)).toBe(true);
  });

  it("routes terminal.status exited to onStatus and stops tracking the session", async () => {
    const h = harness();
    const chunks: Uint8Array[] = [];
    const onStatus = vi.fn();
    await attach(h.ws, h.sockets, "s1", (b) => chunks.push(b), onStatus);

    h.sockets[0].receive({
      kind: "terminal.status",
      payload: { sessionId: "s1", status: "exited", exitCode: 0 },
    });

    expect(onStatus).toHaveBeenCalledWith("exited", 0);
    expect(chunks.some((c) => decode(c).includes("session terminée"))).toBe(true);

    // Session untracked ⇒ a later socket close does not schedule a reconnect.
    h.sockets[0].close();
    expect(h.hasTimer()).toBe(false);
  });

  it("reconnects for a live domain-event subscription even with no terminal (F5)", async () => {
    const h = harness();
    // A live surface subscribes to domain events (no terminal open).
    h.ws.setDomainEventHandler(() => {});
    await h.ws.ensureConnected();
    await vi.waitFor(() => expect(h.ws.getConnectionState()).toBe("connected"));

    // Drop the socket → must schedule a reconnect (the subscription still needs it).
    h.sockets[0].close();
    expect(h.ws.getConnectionState()).toBe("reconnecting");
    expect(h.hasTimer()).toBe(true);

    h.fireTimer();
    await vi.waitFor(() => expect(h.sockets.length).toBe(2));
    await vi.waitFor(() => expect(h.ws.getConnectionState()).toBe("connected"));

    // Events resume flowing to the persisted handler after reconnect.
    const events: unknown[] = [];
    h.ws.setDomainEventHandler((e) => events.push(e));
    h.sockets[1].receive({ kind: "event.domain", payload: { type: "agentBusyChanged", agentId: "a1", busy: true } });
    expect(events).toHaveLength(1);
  });

  it("dispose() moves to closed and clears any pending reconnect", async () => {
    const h = harness();
    await attach(h.ws, h.sockets, "s1", () => {});
    h.sockets[0].close();
    expect(h.ws.getConnectionState()).toBe("reconnecting");
    h.ws.dispose();
    expect(h.ws.getConnectionState()).toBe("closed");
    expect(h.hasTimer()).toBe(false);
  });

  it("disconnect() closes and stops reconnecting but stays reusable (F6 sign-out)", async () => {
    const h = harness();
    h.ws.setDomainEventHandler(() => {});
    await h.ws.ensureConnected();
    await vi.waitFor(() => expect(h.ws.getConnectionState()).toBe("connected"));

    // Sign-out: soft teardown. State is closed, the reconnect loop is dropped, and
    // the detached socket's onclose cannot flip the state back to reconnecting.
    h.ws.disconnect();
    expect(h.ws.getConnectionState()).toBe("closed");
    expect(h.hasTimer()).toBe(false);
    expect(h.ws.needsReconnect()).toBe(false);

    // Reusable: a later ensureConnected reconnects as a fresh socket (re-pairing).
    await h.ws.ensureConnected();
    await vi.waitFor(() => expect(h.sockets.length).toBe(2));
    await vi.waitFor(() => expect(h.ws.getConnectionState()).toBe("connected"));
  });

  it("calls onReconnectFailed when a reconnect attempt fails (auth probe seam)", async () => {
    const failed = vi.fn();
    let openNextSocket = true;
    let timerFn: (() => void) | null = null;
    const sockets: FakeSocket[] = [];
    const ws = new WsLiveClient({
      wsUrl: "wss://host",
      reconnectBaseMs: 1,
      onReconnectFailed: failed,
      socketFactory: () => {
        const s = new FakeSocket();
        sockets.push(s);
        // First socket connects; the reconnect socket errors (upgrade rejected).
        if (openNextSocket) queueMicrotask(() => s.onopen?.());
        else queueMicrotask(() => s.onerror?.(new Error("upgrade rejected")));
        return s;
      },
      setTimeoutImpl: (fn) => {
        timerFn = fn;
        return 1;
      },
      clearTimeoutImpl: () => {
        timerFn = null;
      },
    });

    ws.setDomainEventHandler(() => {});
    await ws.ensureConnected();
    await vi.waitFor(() => expect(ws.getConnectionState()).toBe("connected"));

    // Drop the socket; the next connect attempt will fail the upgrade.
    openNextSocket = false;
    sockets[0].close();
    expect(ws.getConnectionState()).toBe("reconnecting");
    const fire = timerFn as (() => void) | null;
    timerFn = null;
    fire?.();

    // The failed reconnect fires the probe hook and schedules another attempt.
    await vi.waitFor(() => expect(failed).toHaveBeenCalled());
    expect(timerFn).not.toBeNull();
  });
});