arboretum/packages/vscode/src/api/ws-client.ts

// Client WebSocket multiplexé Node : portage du client web (packages/web/src/lib/ws-client.ts).
// UNE connexion, plusieurs canaux ; handshake hello/hello_ok, sub des topics, corrélation FIFO des
// 'attached', flow control par ACK (cumul d'octets PAR CANAL, remis à 0 à chaque RESYNC), reconnexion
// avec backoff. Différences vs web :
//   - transport injecté (RawSocket) → testable sans réseau ; l'implémentation par défaut utilise `ws`
//     avec l'en-tête `Authorization: Bearer` posé à l'upgrade (le navigateur, lui, envoie le cookie).
//   - pas de callback de rendu (xterm) : on ACK les octets AU MOMENT où on les remet au sink (onDidWrite
//     d'un Pseudoterminal émet de façon synchrone). Un ACK traînant débouncé flushe le reliquat < 64 Ko.
import {
  BINARY_FRAME,
  FLOW,
  PROTOCOL_VERSION,
  decodeBinaryFrame,
  type ClientMessage,
  type ServerMessage,
} from '@arboretum/shared';
import { wsUrlFromBase } from '../config.js';

export type WsStatus = 'idle' | 'connecting' | 'open' | 'reconnecting';
export type DetachReason = 'client' | 'session_exit' | 'replaced';

export type SessionEvent = Extract<ServerMessage, { type: 'session_update' | 'session_exit' }>;
export type WorktreeEvent = Extract<ServerMessage, { type: 'repo_update' | 'repo_removed' | 'worktree_update' | 'worktree_removed' }>;
export type GroupEvent = Extract<ServerMessage, { type: 'group_update' | 'group_removed' }>;

/** Sink d'un canal : reçoit les octets BRUTS (le décodage UTF-8 est délégué au SessionBridge). */
export interface AttachmentSink {
  data(payload: Uint8Array): void;
  reset(): void;
  detached(reason: DetachReason): void;
  controlChanged(controlling: boolean): void;
}

export interface AttachOptions {
  sessionId: string;
  mode: 'interactive' | 'observer';
  cols: number;
  rows: number;
  sink: AttachmentSink;
}

/** Transport minimal : abstrait pour les tests ; l'impl par défaut enveloppe `ws`. */
export interface RawSocket {
  send(data: string): void;
  close(): void;
  onOpen(cb: () => void): void;
  onText(cb: (text: string) => void): void;
  onBinary(cb: (data: Uint8Array) => void): void;
  onClose(cb: () => void): void;
  onError(cb: (err: Error) => void): void;
}

export type SocketFactory = (wsUrl: string, token: string) => RawSocket;

const BACKOFF_MIN_MS = 500;
const BACKOFF_MAX_MS = 10_000;

export class Attachment {
  channel = -1;
  controlling = false;
  closed = false;
  cols: number;
  rows: number;
  epoch = 0;
  processedBytes = 0;
  lastAckBytes = 0;
  trailingAckTimer: ReturnType<typeof setTimeout> | null = null;
  pending: { resolve: (a: Attachment) => void; reject: (err: Error) => void } | null = null;

  constructor(
    private readonly client: ArbWsClient,
    readonly sessionId: string,
    readonly mode: 'interactive' | 'observer',
    readonly sink: AttachmentSink,
    cols: number,
    rows: number,
  ) {
    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 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 },
    );
  }

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

export interface ArbWsClientOptions {
  socketFactory: SocketFactory;
  onStatus?: (status: WsStatus) => void;
  onSession?: (e: SessionEvent) => void;
  onWorktree?: (e: WorktreeEvent) => void;
  onGroup?: (e: GroupEvent) => void;
  /** P10 : une session managée terminée vient d'être auto-archivée. */
  onSessionArchived?: (sessionId: string) => void;
  /** appelé au premier hello_ok après (re)connexion → re-seed du store via REST. */
  onReady?: () => void;
  log?: (msg: string) => void;
}

export class ArbWsClient {
  status: WsStatus = 'idle';

  private baseUrl = '';
  private token = '';
  private socket: RawSocket | null = null;
  private ready = false;
  private stopped = true;
  private backoffMs = BACKOFF_MIN_MS;
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  private readonly attachments = new Set<Attachment>();
  private readonly byChannel = new Map<number, Attachment>();
  private awaitingAttached: Attachment[] = [];

  constructor(private readonly opts: ArbWsClientOptions) {}

  /** (Re)configure l'URL/token ; si la connexion était active, elle est relancée. */
  configure(baseUrl: string, token: string): void {
    const wasActive = !this.stopped;
    this.baseUrl = baseUrl;
    this.token = token;
    if (wasActive) {
      this.hardReset();
      this.connect();
    }
  }

  connect(): void {
    this.stopped = false;
    if (this.socket || this.reconnectTimer) return;
    if (!this.token) return; // pas authentifié : on n'ouvre rien
    this.setStatus('connecting');
    this.openSocket();
  }

  disconnect(): void {
    this.stopped = true;
    this.hardReset();
    this.setStatus('idle');
  }

  /** 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);
    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;
  }

  sendControl(msg: ClientMessage): void {
    if (!this.ready || !this.socket) return;
    try {
      this.socket.send(JSON.stringify(msg));
    } catch {
      /* socket en cours de fermeture : la reconnexion rejouera */
    }
  }

  releaseAttachment(att: Attachment): void {
    this.attachments.delete(att);
    if (att.trailingAckTimer) {
      clearTimeout(att.trailingAckTimer);
      att.trailingAckTimer = null;
    }
    if (att.channel >= 0) {
      this.byChannel.delete(att.channel);
      this.sendControl({ type: 'detach', channel: att.channel });
      att.channel = -1;
    }
  }

  // ---- interne ----

  private setStatus(s: WsStatus): void {
    this.status = s;
    this.opts.onStatus?.(s);
  }

  /** ferme la socket courante et invalide tous les états sans toucher au flag stopped. */
  private hardReset(): void {
    this.clearReconnectTimer();
    const socket = this.socket;
    this.socket = null;
    this.ready = false;
    for (const att of this.attachments) {
      if (att.trailingAckTimer) {
        clearTimeout(att.trailingAckTimer);
        att.trailingAckTimer = null;
      }
      att.pending?.reject(new Error('disconnected'));
      att.pending = null;
      att.closed = true;
      att.sink.detached('client');
    }
    this.attachments.clear();
    this.byChannel.clear();
    this.awaitingAttached = [];
    socket?.close();
  }

  private openSocket(): void {
    const socket = this.opts.socketFactory(wsUrlFromBase(this.baseUrl), this.token);
    this.socket = socket;
    socket.onOpen(() => {
      const hello: ClientMessage = { type: 'hello', protocol: PROTOCOL_VERSION };
      socket.send(JSON.stringify(hello));
    });
    socket.onText((text) => this.handleText(text));
    socket.onBinary((data) => this.handleBinary(data));
    socket.onClose(() => this.handleClose(socket));
    socket.onError((err) => this.opts.log?.(`ws error: ${err.message}`));
  }

  private handleClose(socket: RawSocket): 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 (att.trailingAckTimer) {
        clearTimeout(att.trailingAckTimer);
        att.trailingAckTimer = null;
      }
    }
    if (this.stopped) {
      this.setStatus('idle');
      return;
    }
    this.setStatus('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 });
  }

  private handleText(text: string): void {
    let msg: ServerMessage;
    try {
      msg = JSON.parse(text) 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 : reset terminal, replay, compteurs à zéro (le replay ne compte pas)
      att.epoch += 1;
      att.processedBytes = 0;
      att.lastAckBytes = 0;
      att.sink.reset();
      if (frame.payload.byteLength > 0) att.sink.data(frame.payload);
      this.scheduleTrailingAck(att, att.epoch);
      return;
    }
    if (frame.type !== BINARY_FRAME.OUTPUT) return;

    // Émission synchrone vers le sink (onDidWrite) → on ACK directement les octets remis.
    const length = frame.payload.byteLength;
    att.sink.data(frame.payload);
    if (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 });
    }
    // Filet : flushe le reliquat < ACK_EVERY_BYTES pour que le serveur sorte toujours de pause.
    this.scheduleTrailingAck(att, att.epoch);
  }

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

  private handleServerMessage(msg: ServerMessage): void {
    switch (msg.type) {
      case 'hello_ok': {
        this.ready = true;
        this.backoffMs = BACKOFF_MIN_MS;
        this.setStatus('open');
        this.sendControl({ type: 'sub', topics: ['sessions', 'worktrees', 'groups'] });
        for (const att of this.attachments) {
          if (!att.closed) this.sendAttach(att);
        }
        this.opts.onReady?.();
        return;
      }
      case 'attached': {
        const att = this.awaitingAttached.shift();
        if (!att) return;
        if (att.closed) {
          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) {
          att.sink.controlChanged(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.detached(msg.reason);
        return;
      }
      case 'control_changed': {
        const att = this.byChannel.get(msg.channel);
        if (!att || att.closed) return;
        att.controlling = msg.controlling;
        att.sink.controlChanged(msg.controlling);
        return;
      }
      case 'session_update':
      case 'session_exit':
        this.opts.onSession?.(msg);
        return;
      case 'session_archived':
        this.opts.onSessionArchived?.(msg.sessionId);
        return;
      case 'repo_update':
      case 'repo_removed':
      case 'worktree_update':
      case 'worktree_removed':
        this.opts.onWorktree?.(msg);
        return;
      case 'group_update':
      case 'group_removed':
        this.opts.onGroup?.(msg);
        return;
      case 'error': {
        if (msg.channel !== undefined) {
          this.opts.log?.(`channel ${msg.channel}: ${msg.code}: ${msg.message}`);
          return;
        }
        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.detached('session_exit');
          }
          return;
        }
        this.opts.log?.(`${msg.code}: ${msg.message}`);
        return;
      }
      case 'pong':
        return;
    }
  }
}