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

/**
 * WebSocket live client for the web transport — ticket #13, lot F3 (finalised
 * from the F1 skeleton).
 *
 * Owns a single WS connection to `{wsUrl}/ws/live` (authenticated by the session
 * cookie at the upgrade — same-origin, no URL secret) and multiplexes over it,
 * aligned on the **B5 frozen frame contract** (`crates/app-tauri/src/server.rs`):
 *
 *  - client→server: `terminal.open` `{cwd,rows,cols}`, `terminal.attach`
 *    `{sessionId,rows?,cols?,lastSeq?}`, `terminal.input` `{sessionId,bytesBase64}`,
 *    `terminal.resize` `{sessionId,rows,cols}`, `terminal.detach` `{sessionId}`,
 *    `terminal.close` `{sessionId}`, `ping`.
 *  - server→client: `terminal.attached` (ack: `{session,scrollback,nextSeq,status,
 *    gap,assignedConversationId}`), `terminal.output` `{sessionId,seq,bytesBase64}`,
 *    `terminal.status` `{sessionId,status,exitCode}`, `error`, `pong`, plus the
 *    low-frequency `event.domain` stream (domain events).
 *
 * F3 additions over F1: a **connection state machine** (connecting / connected /
 * reconnecting / closed) with automatic reconnection + re-attach and bounded
 * scrollback repaint, and per-session **status routing** (`exited`). The terminal
 * lifecycle is owned here so a browser reload / network drop transparently
 * re-attaches the surviving server-side PTY. Multi-tab "last attach wins" is
 * silent on the wire (the evicted attachment simply stops receiving output — see
 * the F3 report), so there is nothing to crash on; the socket-close path handles
 * disconnection uniformly.
 *
 * Lives in `src/adapters/**`; touches no `@tauri-apps/api`. The port
 * (`TerminalGateway`/`TerminalHandle`) and `TerminalView` are unchanged; the UI
 * indication of connection state is written into xterm through the session's own
 * output sink (a notice line), so no component needs to know about this client.
 */

import type { DomainEvent, Unsubscribe } from "@/domain";
import {
  base64ToBytes,
  bytesToBase64,
  type AttachedPayload,
  type ClientFrame,
  type ClientFrameKind,
  type OutputPayload,
  type ServerFrame,
  type StatusPayload,
} from "./frames";

/** Minimal WebSocket surface used here; injectable so tests pass a fake. */
export interface WebSocketLike {
  send(data: string): void;
  close(): void;
  onopen: (() => void) | null;
  onmessage: ((event: { data: string }) => void) | null;
  onerror: ((event: unknown) => void) | null;
  onclose: (() => void) | null;
}

/** Factory building a {@link WebSocketLike} for a URL (defaults to global WS). */
export type WebSocketFactory = (url: string) => WebSocketLike;

/** Injectable deferred-call primitives (defaults to global timers). */
export type SetTimeoutLike = (fn: () => void, ms: number) => unknown;
export type ClearTimeoutLike = (handle: unknown) => void;

/** UI-facing connection state of the live socket. */
export type ConnectionState = "connecting" | "connected" | "reconnecting" | "closed";

/** Configuration for the live client. */
export interface WsLiveClientConfig {
  /** Base WS URL, e.g. `wss://host:port`. No trailing slash. */
  wsUrl: string;
  /** Bearer token (dev/tests only). Prod auth is the session cookie at upgrade. */
  token?: string;
  /** Injected WebSocket factory (defaults to `new WebSocket(url)`). */
  socketFactory?: WebSocketFactory;
  /** Base reconnect delay in ms (exponential backoff, capped). Default 500. */
  reconnectBaseMs?: number;
  /** Max reconnect delay in ms. Default 10000. */
  reconnectMaxMs?: number;
  /** Injected `setTimeout` (tests drive reconnection deterministically). */
  setTimeoutImpl?: SetTimeoutLike;
  /** Injected `clearTimeout`. */
  clearTimeoutImpl?: ClearTimeoutLike;
  /**
   * Called after a reconnection attempt fails. The browser cannot read the HTTP
   * status of a rejected WS upgrade, so an auth-revoked socket looks like any
   * other outage from here. F6 wires this to a cheap HTTP probe (`health`): a
   * `401` there routes back to pairing (via the invoker's `onUnauthorized`),
   * while a network error just lets the backoff keep retrying.
   */
  onReconnectFailed?: () => void;
}

/** A live terminal subscription tracked for output routing + reconnection replay. */
interface TerminalSubscription {
  /** Delivers raw PTY bytes (and adapter notices) to the xterm view. */
  onData: (bytes: Uint8Array) => void;
  /** Optional lifecycle callback (`exited`), for tests / future UI. */
  onStatus?: (status: string, exitCode?: number | null) => void;
  /** Last output sequence seen — sent as `lastSeq` when re-attaching. */
  lastSeq: number;
}

/** Result of opening/attaching a terminal over the WS. */
export interface TerminalAttachResult {
  sessionId: string;
  /** Bounded scrollback bytes to repaint (empty on a fresh open). */
  scrollback: Uint8Array;
  assignedConversationId?: string;
  status?: string;
}

let frameCounter = 0;
/** Monotonic client frame id (unique per client session). */
function nextFrameId(): string {
  frameCounter += 1;
  return `c${frameCounter}`;
}

const encoder = new TextEncoder();
/** Encodes a human notice line to bytes for writing into xterm. */
function notice(text: string): Uint8Array {
  return encoder.encode(`\r\n\x1b[2m[${text}]\x1b[0m\r\n`);
}

/** Concatenates an attached-frame scrollback list into a single byte buffer. */
function scrollbackBytes(payload: AttachedPayload): Uint8Array {
  const chunks = payload.scrollback.map((c) => base64ToBytes(c.bytesBase64));
  const total = chunks.reduce((n, c) => n + c.length, 0);
  const out = new Uint8Array(total);
  let offset = 0;
  for (const c of chunks) {
    out.set(c, offset);
    offset += c.length;
  }
  return out;
}

/**
 * Manages the single live WebSocket, the connection state machine, and the live
 * terminal sessions (so reconnection re-attaches them transparently).
 */
export class WsLiveClient {
  private readonly wsUrl: string;
  private readonly token?: string;
  private readonly socketFactory: WebSocketFactory;
  private readonly reconnectBaseMs: number;
  private readonly reconnectMaxMs: number;
  private readonly setTimeoutImpl: SetTimeoutLike;
  private readonly clearTimeoutImpl: ClearTimeoutLike;
  private readonly onReconnectFailed?: () => void;

  private socket: WebSocketLike | null = null;
  private opening: Promise<void> | null = null;
  private disposed = false;
  private everConnected = false;
  private connState: ConnectionState = "connecting";
  private reconnectAttempt = 0;
  private reconnectHandle: unknown = null;

  /** Per-session PTY output sinks (sessionId → onData). Low-level routing map. */
  private readonly outputSinks = new Map<string, (bytes: Uint8Array) => void>();
  /** Live terminal subscriptions tracked for reconnection replay + status. */
  private readonly terminals = new Map<string, TerminalSubscription>();
  /** Pending command acknowledgements, keyed by client frame id. */
  private readonly pending = new Map<
    string,
    { resolve: (frame: ServerFrame) => void; reject: (err: unknown) => void }
  >();
  /** The single domain-event handler (set by the system gateway). */
  private domainEventHandler: ((event: DomainEvent) => void) | null = null;
  /** Connection-state listeners (UI / tests). */
  private readonly connListeners = new Set<(state: ConnectionState) => void>();

  constructor(config: WsLiveClientConfig) {
    this.wsUrl = config.wsUrl.replace(/\/+$/, "");
    this.token = config.token;
    this.socketFactory =
      config.socketFactory ??
      ((url: string) => new WebSocket(url) as unknown as WebSocketLike);
    this.reconnectBaseMs = config.reconnectBaseMs ?? 500;
    this.reconnectMaxMs = config.reconnectMaxMs ?? 10000;
    this.setTimeoutImpl =
      config.setTimeoutImpl ?? ((fn, ms) => setTimeout(fn, ms));
    this.clearTimeoutImpl =
      config.clearTimeoutImpl ?? ((h) => clearTimeout(h as ReturnType<typeof setTimeout>));
    this.onReconnectFailed = config.onReconnectFailed;
  }

  // -------------------------------------------------------------------------
  // Connection state
  // -------------------------------------------------------------------------

  /** Current connection state. */
  getConnectionState(): ConnectionState {
    return this.connState;
  }

  /** Subscribes to connection-state transitions (returns an unsubscribe). */
  onConnectionStateChange(listener: (state: ConnectionState) => void): Unsubscribe {
    this.connListeners.add(listener);
    return () => this.connListeners.delete(listener);
  }

  private setConnState(state: ConnectionState): void {
    if (this.connState === state) return;
    this.connState = state;
    for (const listener of this.connListeners) listener(state);
  }

  // -------------------------------------------------------------------------
  // Domain events (system gateway)
  // -------------------------------------------------------------------------

  setDomainEventHandler(handler: (event: DomainEvent) => void): void {
    this.domainEventHandler = handler;
  }
  clearDomainEventHandler(): void {
    this.domainEventHandler = null;
  }

  // -------------------------------------------------------------------------
  // Low-level output sink map (used by the agent gateway too)
  // -------------------------------------------------------------------------

  setOutputSink(sessionId: string, onData: (bytes: Uint8Array) => void): void {
    this.outputSinks.set(sessionId, onData);
  }
  removeOutputSink(sessionId: string): void {
    this.outputSinks.delete(sessionId);
  }

  // -------------------------------------------------------------------------
  // Connection lifecycle
  // -------------------------------------------------------------------------

  /** Ensures the socket is connected, connecting on first use. */
  async ensureConnected(): Promise<void> {
    if (this.socket) return;
    if (this.opening) return this.opening;
    return this.connect();
  }

  private connect(): Promise<void> {
    this.setConnState(this.everConnected ? "reconnecting" : "connecting");
    this.opening = new Promise<void>((resolve, reject) => {
      // Prod: the session cookie rides the upgrade (same-origin). The optional
      // token is a dev/test convenience only (never a prod secret in the URL).
      const url = this.token
        ? `${this.wsUrl}/ws/live?token=${encodeURIComponent(this.token)}`
        : `${this.wsUrl}/ws/live`;
      const socket = this.socketFactory(url);
      socket.onopen = () => {
        this.socket = socket;
        this.opening = null;
        this.everConnected = true;
        this.reconnectAttempt = 0;
        this.setConnState("connected");
        resolve();
      };
      socket.onerror = (event) => {
        this.opening = null;
        reject(event);
      };
      socket.onclose = () => this.handleSocketClosed();
      socket.onmessage = (event) => this.handleMessage(event.data);
    });
    return this.opening;
  }

  private handleSocketClosed(): void {
    this.socket = null;
    this.opening = null;
    this.rejectAllPending({ code: "WS_CLOSED", message: "socket closed" });
    if (this.disposed) {
      this.setConnState("closed");
      return;
    }
    this.setConnState("reconnecting");
    // Reconnect while anything still needs the socket: a live terminal to
    // re-attach (F3) OR an active domain-event subscription (F5 live surfaces).
    // Otherwise stay idle until the next explicit use.
    if (this.terminals.size > 0) {
      for (const sub of this.terminals.values()) sub.onData(notice("déconnecté — reconnexion…"));
    }
    if (this.needsReconnect()) this.scheduleReconnect();
  }

  /** Whether the socket should auto-reconnect (a terminal or a live subscription). */
  needsReconnect(): boolean {
    return this.terminals.size > 0 || this.domainEventHandler !== null;
  }

  private scheduleReconnect(): void {
    if (this.disposed || this.reconnectHandle) return;
    const delay = Math.min(
      this.reconnectMaxMs,
      this.reconnectBaseMs * 2 ** this.reconnectAttempt,
    );
    this.reconnectAttempt += 1;
    this.reconnectHandle = this.setTimeoutImpl(() => {
      this.reconnectHandle = null;
      void this.reconnectNow();
    }, delay);
  }

  private async reconnectNow(): Promise<void> {
    if (this.disposed) return;
    try {
      await this.connect();
      await this.reattachAll();
    } catch {
      // Still down: probe (an auth-revoked upgrade is indistinguishable from an
      // outage here — the probe surfaces a 401 as a pairing bounce), then back
      // off and retry (unless nothing needs the socket anymore).
      this.onReconnectFailed?.();
      if (this.needsReconnect()) this.scheduleReconnect();
    }
  }

  /** Re-attaches every tracked terminal after a reconnect and repaints scrollback. */
  private async reattachAll(): Promise<void> {
    for (const [sessionId, sub] of this.terminals) {
      try {
        const ack = await this.send("terminal.attach", {
          sessionId,
          lastSeq: sub.lastSeq,
        });
        const payload = ack.payload as unknown as AttachedPayload;
        sub.onData(notice("reconnecté"));
        const bytes = scrollbackBytes(payload);
        if (bytes.length > 0) sub.onData(bytes);
        if (typeof payload.nextSeq === "number") {
          sub.lastSeq = Math.max(0, payload.nextSeq - 1);
        }
      } catch {
        // A single session failing to re-attach (e.g. server killed it) must not
        // block the others; leave it tracked for the next reconnect cycle.
      }
    }
  }

  // -------------------------------------------------------------------------
  // Terminal sessions (high-level; used by HttpTerminalGateway)
  // -------------------------------------------------------------------------

  /** Opens a fresh terminal (`terminal.open`) and starts tracking it. */
  openTerminal(params: {
    cwd: string;
    rows: number;
    cols: number;
    nodeId?: string | null;
    onData: (bytes: Uint8Array) => void;
    onStatus?: (status: string, exitCode?: number | null) => void;
  }): Promise<TerminalAttachResult> {
    return this.attachInternal(
      "terminal.open",
      { cwd: params.cwd, rows: params.rows, cols: params.cols, nodeId: params.nodeId ?? null },
      { onData: params.onData, onStatus: params.onStatus, lastSeq: 0 },
    );
  }

  /**
   * Launches a CLI agent (`agent.launch`) over the same PTY channel and tracks
   * it exactly like a terminal (B6): the ack is a unified `terminal.attached`
   * carrying `sessionId` + `assignedConversationId`, then output/input/resize and
   * reconnection replay reuse the terminal mechanics. Structured agents are
   * refused server-side with `UNSUPPORTED` (this channel is PTY-only); the
   * rejected `send` surfaces that error to the caller unchanged.
   */
  launchAgent(params: {
    projectId: string;
    agentId: string;
    rows: number;
    cols: number;
    nodeId?: string | null;
    conversationId?: string | null;
    onData: (bytes: Uint8Array) => void;
    onStatus?: (status: string, exitCode?: number | null) => void;
  }): Promise<TerminalAttachResult> {
    return this.attachInternal(
      "agent.launch",
      {
        projectId: params.projectId,
        agentId: params.agentId,
        nodeId: params.nodeId ?? null,
        rows: params.rows,
        cols: params.cols,
        conversationId: params.conversationId ?? null,
      },
      { onData: params.onData, onStatus: params.onStatus, lastSeq: 0 },
    );
  }

  /** Re-attaches to an existing terminal (`terminal.attach`) and tracks it. */
  attachTerminal(params: {
    sessionId: string;
    rows?: number;
    cols?: number;
    lastSeq?: number;
    onData: (bytes: Uint8Array) => void;
    onStatus?: (status: string, exitCode?: number | null) => void;
  }): Promise<TerminalAttachResult> {
    return this.attachInternal(
      "terminal.attach",
      {
        sessionId: params.sessionId,
        rows: params.rows,
        cols: params.cols,
        lastSeq: params.lastSeq ?? null,
      },
      { onData: params.onData, onStatus: params.onStatus, lastSeq: params.lastSeq ?? 0 },
      params.sessionId,
    );
  }

  private async attachInternal(
    kind: ClientFrameKind,
    payload: Record<string, unknown>,
    sub: TerminalSubscription,
    knownSessionId?: string,
  ): Promise<TerminalAttachResult> {
    const ack = await this.send(kind, payload);
    const ap = ack.payload as unknown as AttachedPayload;
    const sessionId = knownSessionId ?? ap.session.sessionId;
    if (typeof ap.nextSeq === "number") sub.lastSeq = Math.max(0, ap.nextSeq - 1);
    this.terminals.set(sessionId, sub);
    // Route output to the subscription's onData (the low-level map is the single
    // source of routing; lastSeq is tracked in handleMessage).
    this.setOutputSink(sessionId, sub.onData);
    return {
      sessionId,
      scrollback: scrollbackBytes(ap),
      assignedConversationId: ap.assignedConversationId,
      status: ap.status,
    };
  }

  /** Detaches the view (server PTY keeps running): stop tracking + `terminal.detach`. */
  async detachTerminal(sessionId: string): Promise<void> {
    this.removeTerminal(sessionId);
    await this.sendFireAndForget("terminal.detach", { sessionId });
  }

  /** Kills the PTY (`terminal.close`) and stops tracking. */
  async closeTerminalSession(sessionId: string): Promise<void> {
    this.removeTerminal(sessionId);
    await this.send("terminal.close", { sessionId });
  }

  private removeTerminal(sessionId: string): void {
    this.terminals.delete(sessionId);
    this.outputSinks.delete(sessionId);
  }

  // -------------------------------------------------------------------------
  // Frame I/O
  // -------------------------------------------------------------------------

  /**
   * Sends a client frame and resolves with its acknowledgement (the server frame
   * whose `replyTo` equals this frame's `id`).
   */
  async send(
    kind: ClientFrameKind,
    payload: Record<string, unknown>,
  ): Promise<ServerFrame> {
    await this.ensureConnected();
    const socket = this.socket;
    if (!socket) {
      throw { code: "WS_CLOSED", message: "socket not connected" };
    }
    const id = nextFrameId();
    const frame: ClientFrame = { id, kind, payload };
    const ack = new Promise<ServerFrame>((resolve, reject) => {
      this.pending.set(id, { resolve, reject });
    });
    socket.send(JSON.stringify(frame));
    return ack;
  }

  /** Sends a fire-and-forget frame (no acknowledgement awaited). */
  async sendFireAndForget(
    kind: ClientFrameKind,
    payload: Record<string, unknown>,
  ): Promise<void> {
    await this.ensureConnected();
    const id = nextFrameId();
    const frame: ClientFrame = { id, kind, payload };
    this.socket?.send(JSON.stringify(frame));
  }

  /** Convenience: base64-encode bytes for an `input` frame. */
  static encodeInput(bytes: Uint8Array): string {
    return bytesToBase64(bytes);
  }

  /**
   * Soft teardown for a sign-out / auth bounce (F6): closes the socket, drops all
   * live state and stops the reconnect loop, but — unlike {@link dispose} — leaves
   * the client reusable. A later `ensureConnected` (after re-pairing) reconnects
   * as a fresh `connecting`. The socket's own `onclose` is detached first so it
   * cannot flip the state back to `reconnecting`.
   */
  disconnect(): void {
    if (this.reconnectHandle) {
      this.clearTimeoutImpl(this.reconnectHandle);
      this.reconnectHandle = null;
    }
    this.rejectAllPending({ code: "WS_CLOSED", message: "client disconnected" });
    this.outputSinks.clear();
    this.terminals.clear();
    this.domainEventHandler = null;
    const socket = this.socket;
    if (socket) {
      socket.onopen = null;
      socket.onmessage = null;
      socket.onerror = null;
      socket.onclose = null;
      socket.close();
    }
    this.socket = null;
    this.opening = null;
    this.everConnected = false;
    this.reconnectAttempt = 0;
    this.setConnState("closed");
  }

  /** Closes the socket and rejects everything pending. */
  dispose(): void {
    this.disposed = true;
    if (this.reconnectHandle) {
      this.clearTimeoutImpl(this.reconnectHandle);
      this.reconnectHandle = null;
    }
    this.rejectAllPending({ code: "WS_CLOSED", message: "client disposed" });
    this.outputSinks.clear();
    this.terminals.clear();
    this.domainEventHandler = null;
    this.socket?.close();
    this.socket = null;
    this.opening = null;
    this.setConnState("closed");
  }

  private rejectAllPending(err: unknown): void {
    for (const { reject } of this.pending.values()) reject(err);
    this.pending.clear();
  }

  private handleMessage(data: string): void {
    let frame: ServerFrame;
    try {
      frame = JSON.parse(data) as ServerFrame;
    } catch {
      return; // Ignore malformed frames.
    }

    // Route unsolicited streams first (no replyTo).
    switch (frame.kind) {
      case "terminal.output": {
        const payload = frame.payload as unknown as OutputPayload;
        const sub = this.terminals.get(payload.sessionId);
        if (sub && typeof payload.seq === "number") sub.lastSeq = payload.seq;
        const sink = this.outputSinks.get(payload.sessionId);
        if (sink) sink(base64ToBytes(payload.bytesBase64));
        return;
      }
      case "terminal.status": {
        const payload = frame.payload as unknown as StatusPayload;
        this.handleStatus(payload);
        return;
      }
      case "event.domain": {
        this.domainEventHandler?.(frame.payload as unknown as DomainEvent);
        return;
      }
      default:
        break;
    }

    // Otherwise it is a reply/ack correlated by `replyTo`.
    if (frame.replyTo) {
      const waiter = this.pending.get(frame.replyTo);
      if (waiter) {
        this.pending.delete(frame.replyTo);
        if (frame.kind === "error") waiter.reject(frame.payload);
        else waiter.resolve(frame);
      }
    }
  }

  private handleStatus(payload: StatusPayload): void {
    const sub = this.terminals.get(payload.sessionId);
    if (!sub) return;
    if (payload.status === "exited" || payload.status === "closed") {
      const code = payload.exitCode;
      sub.onData(
        notice(
          code === null || code === undefined
            ? "session terminée"
            : `session terminée (code ${code})`,
        ),
      );
      sub.onStatus?.(payload.status, payload.exitCode);
      // The PTY is gone: stop tracking so a later disconnect does not try to
      // re-attach a dead session.
      this.removeTerminal(payload.sessionId);
    } else {
      sub.onStatus?.(payload.status, payload.exitCode);
    }
  }
}