arboretum/packages/web/src/lib/ws-client.ts

// Client WebSocket multiplexé : UNE seule connexion pour toute l'app.
// Handshake hello/hello_ok, reconnexion avec backoff (ré-attache des terminaux
// ouverts + re-sub), corrélation FIFO des 'attached', flow control par ACK
// (cumul d'octets traités PAR CANAL, remis à zéro à chaque RESYNC).
import { ref, type Ref } from 'vue';
import {
  BINARY_FRAME,
  FLOW,
  PROTOCOL_VERSION,
  decodeBinaryFrame,
  type ClientMessage,
  type ServerMessage,
} from '@arboretum/shared';

export type WsStatus = 'idle' | 'connecting' | 'open' | 'reconnecting';

export type DetachReason = 'client' | 'session_exit' | 'replaced';

/** Interface minimale attendue côté terminal (implémentée par xterm). */
export interface TerminalSink {
  /** le callback est invoqué quand xterm a réellement traité le chunk (base du flow control) */
  write(data: Uint8Array, callback?: () => void): void;
  /** reset complet avant le replay d'une frame RESYNC */
  reset(): void;
  onDetached(reason: DetachReason): void;
  onControlChanged(controlling: boolean): void;
  /**
   * Erreur serveur portant SUR ce canal (frappe refusée, canal inconnu). Optionnel, mais sans lui
   * ces refus se perdaient dans un `console.warn` : l'utilisateur tapait et rien ne se passait.
   */
  onChannelError?(code: string, message: string): void;
}

export type SessionEvent = Extract<ServerMessage, { type: 'session_update' | 'session_exit' | 'session_archived' }>;
export type WorktreeEvent = Extract<ServerMessage, { type: 'repo_update' | 'repo_removed' | 'worktree_update' | 'worktree_removed' }>;
export type GroupEvent = Extract<ServerMessage, { type: 'group_update' | 'group_removed' }>;
/** P11 : un réglage a changé (diffusé au topic 'settings'). */
export type SettingsUpdateEvent = Extract<ServerMessage, { type: 'settings_update' }>;
/** P12 : progression/fin d'un clone (diffusée au topic 'clones'). */
export type CloneEvent = Extract<ServerMessage, { type: 'clone_update' }>;
/** P7 : signal ciblé « le détail (changes/diff) d'un worktree regardé a changé ». */
export type WorktreeChangesEvent = Extract<ServerMessage, { type: 'worktree_changes' }>;

export interface AttachOptions {
  sessionId: string;
  mode: 'interactive' | 'observer';
  cols: number;
  rows: number;
  sink: TerminalSink;
  /** false = attache d'écriture seule, sans affichage (cf. le champ `screen` du protocole). */
  screen?: boolean;
}

const BACKOFF_MIN_MS = 500;
const BACKOFF_MAX_MS = 10_000;

export class Attachment {
  /** canal courant : réassigné à chaque reconnexion, -1 tant que non attaché */
  channel = -1;
  controlling = false;
  closed = false;
  cols: number;
  rows: number;
  /** époque de resync : invalide les callbacks d'écritures encore en vol après un reset */
  epoch = 0;
  /** octets de payload réellement traités par le terminal depuis le dernier resync */
  processedBytes = 0;
  /** dernier cumul envoyé en ack (repart à 0 à chaque resync, comme côté serveur) */
  lastAckBytes = 0;
  /** ACK traînant : garantit l'envoi du reliquat même quand le flux s'arrête (anti-deadlock) */
  trailingAckTimer: ReturnType<typeof setTimeout> | null = null;
  /** résolution du premier 'attached' de ce canal (corrélation FIFO) */
  pending: { resolve: (a: Attachment) => void; reject: (err: Error) => void } | null = null;

  constructor(
    private readonly client: WsClient,
    readonly sessionId: string,
    readonly mode: 'interactive' | 'observer',
    readonly sink: TerminalSink,
    cols: number,
    rows: number,
    readonly screen = true,
  ) {
    this.cols = cols;
    this.rows = rows;
  }

  sendStdin(data: string): void {
    if (this.closed || this.channel < 0) return;
    this.client.sendControl({ type: 'stdin', channel: this.channel, data });
  }

  /** Répond à un dialogue Claude (P4-A) sans clavier ; le serveur traduit en keystrokes. */
  answer(action: 'select' | 'confirm' | 'deny', optionN?: number): void {
    if (this.closed || this.channel < 0) return;
    this.client.sendControl(
      action === 'select' && optionN !== undefined
        ? { type: 'answer', channel: this.channel, action, optionN }
        : { type: 'answer', channel: this.channel, action },
    );
  }

  /** mémorise toujours la taille (réutilisée à la ré-attache) ; n'envoie que si controlling */
  resize(cols: number, rows: number): void {
    this.cols = cols;
    this.rows = rows;
    if (this.closed || this.channel < 0 || !this.controlling) return;
    this.client.sendControl({ type: 'resize', channel: this.channel, cols, rows });
  }

  detach(): void {
    if (this.closed) return;
    this.closed = true;
    this.pending = null;
    this.client.releaseAttachment(this);
  }

  /** Flush ACK ciblé : relance immédiatement le flux si le serveur avait mis ce PTY en pause
   *  (ex. retour-visible de la cellule, dont le rendu était gelé hors-écran). N'ACK que du traité. */
  flushAck(): void {
    this.client.flushAck(this);
  }
}

export class WsClient {
  readonly status: Ref<WsStatus> = ref('idle');

  private socket: WebSocket | null = null;
  private ready = false;
  private stopped = true;
  private visibilityBound = false;
  private backoffMs = BACKOFF_MIN_MS;
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  /** terminaux ouverts, ré-attachés à chaque reconnexion (ordre d'insertion = ordre FIFO) */
  private readonly attachments = new Set<Attachment>();
  private readonly byChannel = new Map<number, Attachment>();
  /** FIFO des attaches en attente : le serveur répond aux 'attach' dans l'ordre */
  private awaitingAttached: Attachment[] = [];
  private readonly sessionListeners = new Set<(e: SessionEvent) => void>();
  private readonly worktreeListeners = new Set<(e: WorktreeEvent) => void>();
  private readonly groupListeners = new Set<(e: GroupEvent) => void>();
  private readonly settingsListeners = new Set<(e: SettingsUpdateEvent) => void>();
  private readonly cloneListeners = new Set<(e: CloneEvent) => void>();
  /** P7 : abonnements ciblés au détail d'un worktree (clé repoId\0path → listeners). */
  private readonly worktreeChangesListeners = new Map<string, Set<(e: WorktreeChangesEvent) => void>>();

  connect(): void {
    this.stopped = false;
    // Au retour en avant-plan, le rendu se débloque : on flushe les ACK pour ne pas attendre le
    // débounce si le serveur s'était mis en pause pendant que l'onglet/PWA était caché.
    if (!this.visibilityBound && typeof document !== 'undefined') {
      this.visibilityBound = true;
      document.addEventListener('visibilitychange', () => {
        if (document.visibilityState === 'visible') this.flushAcks();
      });
    }
    if (this.socket || this.reconnectTimer) return;
    if (this.status.value === 'idle') this.status.value = 'connecting';
    this.openSocket();
  }

  disconnect(): void {
    this.stopped = true;
    this.clearReconnectTimer();
    const socket = this.socket;
    this.socket = null;
    this.ready = false;
    this.status.value = 'idle';
    for (const att of this.attachments) {
      att.closed = true;
      att.pending?.reject(new Error('disconnected'));
      att.pending = null;
      att.sink.onDetached('client');
    }
    this.attachments.clear();
    this.byChannel.clear();
    this.awaitingAttached = [];
    socket?.close();
  }

  /** résout sur le prochain 'attached' corrélé ; rejette si la session est introuvable/terminée */
  attach(opts: AttachOptions): Promise<Attachment> {
    const att = new Attachment(this, opts.sessionId, opts.mode, opts.sink, opts.cols, opts.rows, opts.screen ?? true);
    const promise = new Promise<Attachment>((resolve, reject) => {
      att.pending = { resolve, reject };
    });
    this.attachments.add(att);
    this.connect();
    if (this.ready) this.sendAttach(att);
    return promise;
  }

  /** topics actifs = union des abonnements courants (une seule connexion partagée). */
  private activeTopics(): Array<'sessions' | 'worktrees' | 'groups' | 'settings' | 'clones'> {
    const t: Array<'sessions' | 'worktrees' | 'groups' | 'settings' | 'clones'> = [];
    if (this.sessionListeners.size > 0) t.push('sessions');
    if (this.worktreeListeners.size > 0) t.push('worktrees');
    if (this.groupListeners.size > 0) t.push('groups');
    if (this.settingsListeners.size > 0) t.push('settings');
    if (this.cloneListeners.size > 0) t.push('clones');
    return t;
  }

  private sendSub(): void {
    this.sendControl({ type: 'sub', topics: this.activeTopics() });
  }

  subscribeSessions(listener: (e: SessionEvent) => void): () => void {
    this.sessionListeners.add(listener);
    this.connect();
    this.sendSub();
    return () => {
      this.sessionListeners.delete(listener);
      this.sendSub();
    };
  }

  subscribeWorktrees(listener: (e: WorktreeEvent) => void): () => void {
    this.worktreeListeners.add(listener);
    this.connect();
    this.sendSub();
    return () => {
      this.worktreeListeners.delete(listener);
      this.sendSub();
    };
  }

  subscribeGroups(listener: (e: GroupEvent) => void): () => void {
    this.groupListeners.add(listener);
    this.connect();
    this.sendSub();
    return () => {
      this.groupListeners.delete(listener);
      this.sendSub();
    };
  }

  subscribeSettings(listener: (e: SettingsUpdateEvent) => void): () => void {
    this.settingsListeners.add(listener);
    this.connect();
    this.sendSub();
    return () => {
      this.settingsListeners.delete(listener);
      this.sendSub();
    };
  }

  subscribeClones(listener: (e: CloneEvent) => void): () => void {
    this.cloneListeners.add(listener);
    this.connect();
    this.sendSub();
    return () => {
      this.cloneListeners.delete(listener);
      this.sendSub();
    };
  }

  /**
   * P7 : observe le détail (changes/diff) d'un worktree précis : envoie `watch`, route les
   * `worktree_changes` correspondants vers `listener`, et ré-arme automatiquement après reconnexion.
   * La désinscription envoie `unwatch` quand plus aucun listener ne regarde ce worktree.
   */
  watchWorktree(repoId: string, path: string, listener: (e: WorktreeChangesEvent) => void): () => void {
    const key = `${repoId}\0${path}`;
    let set = this.worktreeChangesListeners.get(key);
    if (!set) {
      set = new Set();
      this.worktreeChangesListeners.set(key, set);
      this.connect();
      this.sendControl({ type: 'watch', repoId, path });
    }
    set.add(listener);
    return () => {
      const s = this.worktreeChangesListeners.get(key);
      if (!s) return;
      s.delete(listener);
      if (s.size === 0) {
        this.worktreeChangesListeners.delete(key);
        this.sendControl({ type: 'unwatch', repoId, path });
      }
    };
  }

  sendControl(msg: ClientMessage): void {
    if (!this.ready || this.socket?.readyState !== WebSocket.OPEN) return;
    this.socket.send(JSON.stringify(msg));
  }

  /** détache côté serveur et oublie l'attache (appelé par Attachment.detach) */
  releaseAttachment(att: Attachment): void {
    this.attachments.delete(att);
    if (att.channel >= 0) {
      this.byChannel.delete(att.channel);
      this.sendControl({ type: 'detach', channel: att.channel });
      att.channel = -1;
    }
    // si l'attache attend encore son 'attached', elle reste dans la FIFO pour
    // préserver la corrélation : closed=true la fera détacher à la réception
  }

  // ---- interne ----

  private openSocket(): void {
    const socket = new WebSocket(`${location.origin.replace(/^http/, 'ws')}/ws`);
    socket.binaryType = 'arraybuffer';
    this.socket = socket;
    socket.onopen = () => {
      const hello: ClientMessage = { type: 'hello', protocol: PROTOCOL_VERSION };
      socket.send(JSON.stringify(hello));
    };
    socket.onmessage = (ev: MessageEvent) => this.handleMessage(ev);
    socket.onclose = () => this.handleClose(socket);
    socket.onerror = () => {
      // onclose suit systématiquement : la reconnexion se joue là-bas
    };
  }

  private handleClose(socket: WebSocket): void {
    if (this.socket !== socket) return;
    this.socket = null;
    this.ready = false;
    this.byChannel.clear();
    this.awaitingAttached = [];
    for (const att of this.attachments) {
      att.channel = -1;
      att.epoch += 1;
    }
    if (this.stopped) {
      this.status.value = 'idle';
      return;
    }
    this.status.value = 'reconnecting';
    this.scheduleReconnect();
  }

  private scheduleReconnect(): void {
    this.clearReconnectTimer();
    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      if (!this.stopped) this.openSocket();
    }, this.backoffMs);
    this.backoffMs = Math.min(this.backoffMs * 2, BACKOFF_MAX_MS);
  }

  private clearReconnectTimer(): void {
    if (this.reconnectTimer !== null) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }
  }

  private sendAttach(att: Attachment): void {
    this.awaitingAttached.push(att);
    this.sendControl({ type: 'attach', sessionId: att.sessionId, mode: att.mode, cols: att.cols, rows: att.rows, screen: att.screen });
  }

  private handleMessage(ev: MessageEvent): void {
    if (ev.data instanceof ArrayBuffer) {
      this.handleBinary(new Uint8Array(ev.data));
      return;
    }
    let msg: ServerMessage;
    try {
      msg = JSON.parse(String(ev.data)) as ServerMessage;
    } catch {
      return;
    }
    this.handleServerMessage(msg);
  }

  private handleBinary(data: Uint8Array): void {
    if (data.byteLength < BINARY_FRAME.HEADER_BYTES) return;
    const frame = decodeBinaryFrame(data);
    const att = this.byChannel.get(frame.channel);
    if (!att || att.closed) return;

    if (frame.type === BINARY_FRAME.RESYNC) {
      // le serveur repart de sentBytes=0 après un resync : reset terminal, replay
      // du ring, et compteurs d'octets remis à zéro (le payload du replay ne compte pas)
      att.epoch += 1;
      att.processedBytes = 0;
      att.lastAckBytes = 0;
      att.sink.reset();
      if (frame.payload.byteLength > 0) att.sink.write(frame.payload);
      // filet anti-stall : le replay ne passe PAS de callback de write → on arme quand même l'ACK
      // traînant pour que le flux suivant ne puisse pas rester gelé (serveur en pause).
      this.scheduleTrailingAck(att, att.epoch);
      return;
    }
    if (frame.type !== BINARY_FRAME.OUTPUT) return;

    const epoch = att.epoch;
    const length = frame.payload.byteLength;
    att.sink.write(frame.payload, () => {
      // compté seulement si aucun resync/reconnexion n'est passé entre temps
      if (att.closed || att.epoch !== epoch || att.channel < 0) return;
      att.processedBytes += length;
      if (att.processedBytes - att.lastAckBytes >= FLOW.ACK_EVERY_BYTES) {
        att.lastAckBytes = att.processedBytes;
        this.sendControl({ type: 'ack', channel: att.channel, bytes: att.processedBytes });
      }
      this.scheduleTrailingAck(att, epoch);
    });
    // Filet anti-stall, INDÉPENDANT du callback de xterm.write : ce callback est throttlé quand
    // l'onglet/PWA passe en arrière-plan (rAF + setTimeout ralentis) → sans cela l'ACK ne partait
    // plus, le serveur dépassait HIGH_WATERMARK et mettait le PTY en pause, et la sortie « chargeait
    // par à-coups » jusqu'à ce que l'utilisateur bouge. On arme donc l'ACK traînant dès la réception.
    this.scheduleTrailingAck(att, epoch);
  }

  /**
   * (Ré)arme l'ACK traînant débouncé : si rien d'autre n'arrive, il envoie le reliquat d'octets
   * réellement traités par xterm (`processedBytes`) pour que le serveur sorte toujours de pause.
   * Appelé à la fois dans le callback de write ET à la réception d'une frame (le callback peut être
   * indéfiniment retardé quand le rendu est gelé). Sémantique inchangée : on n'ACK que du traité.
   */
  private scheduleTrailingAck(att: Attachment, epoch: number): void {
    if (att.trailingAckTimer) clearTimeout(att.trailingAckTimer);
    att.trailingAckTimer = setTimeout(() => {
      att.trailingAckTimer = null;
      if (att.closed || att.epoch !== epoch || att.channel < 0) return;
      if (att.processedBytes > att.lastAckBytes) {
        att.lastAckBytes = att.processedBytes;
        this.sendControl({ type: 'ack', channel: att.channel, bytes: att.processedBytes });
      }
    }, 200);
  }

  /** ACK immédiat du reliquat traité pour tous les terminaux ouverts : au retour en avant-plan,
   *  pour rattraper sans attendre le débounce un serveur éventuellement en pause. */
  private flushAcks(): void {
    for (const att of this.attachments) this.flushAttachmentAck(att);
  }

  /** ACK immédiat du reliquat traité d'UN attachment : n'ACK que du réellement traité (invariant
   *  préservé). Factorisé : utilisé par flushAcks() et par le flush ciblé flushAck(). */
  private flushAttachmentAck(att: Attachment): void {
    if (att.closed || att.channel < 0) return;
    if (att.processedBytes > att.lastAckBytes) {
      att.lastAckBytes = att.processedBytes;
      this.sendControl({ type: 'ack', channel: att.channel, bytes: att.processedBytes });
    }
  }

  /** Flush ACK ciblé d'un seul attachment : relance un PTY mis en pause par le serveur quand la
   *  cellule était cachée (cf. Attachment.flushAck, appelé au retour-visible). */
  flushAck(att: Attachment): void {
    this.flushAttachmentAck(att);
  }

  private handleServerMessage(msg: ServerMessage): void {
    switch (msg.type) {
      case 'hello_ok': {
        this.ready = true;
        this.backoffMs = BACKOFF_MIN_MS;
        this.status.value = 'open';
        this.sendSub();
        // P7 : ré-arme les watchers de worktree (clé repoId\0path) après reconnexion.
        for (const key of this.worktreeChangesListeners.keys()) {
          const sep = key.indexOf('\0');
          this.sendControl({ type: 'watch', repoId: key.slice(0, sep), path: key.slice(sep + 1) });
        }
        // ré-attache des terminaux encore ouverts (nouveaux canaux côté serveur)
        for (const att of this.attachments) {
          if (!att.closed) this.sendAttach(att);
        }
        return;
      }
      case 'attached': {
        const att = this.awaitingAttached.shift();
        if (!att) return;
        if (att.closed) {
          // détaché pendant l'attente : on libère le canal côté serveur
          this.sendControl({ type: 'detach', channel: msg.channel });
          return;
        }
        const wasControlling = att.controlling;
        att.channel = msg.channel;
        att.controlling = msg.controlling;
        this.byChannel.set(msg.channel, att);
        if (att.pending) {
          att.pending.resolve(att);
          att.pending = null;
        } else if (wasControlling !== msg.controlling) {
          // ré-attache après reconnexion : le statut controlling peut avoir changé
          att.sink.onControlChanged(msg.controlling);
        }
        return;
      }
      case 'detached': {
        const att = this.byChannel.get(msg.channel);
        if (!att) return;
        this.byChannel.delete(msg.channel);
        att.channel = -1;
        if (msg.reason === 'client') return;
        att.closed = true;
        this.attachments.delete(att);
        att.sink.onDetached(msg.reason);
        return;
      }
      case 'control_changed': {
        const att = this.byChannel.get(msg.channel);
        if (!att || att.closed) return;
        att.controlling = msg.controlling;
        att.sink.onControlChanged(msg.controlling);
        return;
      }
      case 'session_update':
      case 'session_exit':
      case 'session_archived': {
        for (const cb of this.sessionListeners) cb(msg);
        return;
      }
      case 'repo_update':
      case 'repo_removed':
      case 'worktree_update':
      case 'worktree_removed': {
        for (const cb of this.worktreeListeners) cb(msg);
        return;
      }
      case 'group_update':
      case 'group_removed': {
        for (const cb of this.groupListeners) cb(msg);
        return;
      }
      case 'settings_update': {
        for (const cb of this.settingsListeners) cb(msg);
        return;
      }
      case 'clone_update': {
        for (const cb of this.cloneListeners) cb(msg);
        return;
      }
      case 'worktree_changes': {
        const set = this.worktreeChangesListeners.get(`${msg.repoId}\0${msg.path}`);
        if (set) for (const cb of set) cb(msg);
        return;
      }
      case 'error': {
        if (msg.channel !== undefined) {
          console.warn(`[ws] channel ${msg.channel}: ${msg.code}: ${msg.message}`);
          // Remonté au terminal concerné : une frappe refusée doit se voir, pas seulement en console.
          this.byChannel.get(msg.channel)?.sink.onChannelError?.(msg.code, msg.message);
          return;
        }
        // un échec d'attach est la seule erreur sans canal corrélable à une requête
        if ((msg.code === 'NOT_FOUND' || msg.code === 'SESSION_EXITED') && this.awaitingAttached.length > 0) {
          const att = this.awaitingAttached.shift();
          if (!att) return;
          this.attachments.delete(att);
          if (att.pending) {
            att.pending.reject(new Error(msg.code));
            att.pending = null;
          } else if (!att.closed) {
            att.closed = true;
            att.sink.onDetached('session_exit');
          }
          return;
        }
        console.warn(`[ws] ${msg.code}: ${msg.message}`);
        return;
      }
      case 'pong':
        return;
    }
  }
}

export const wsClient = new WsClient();