arboretum/packages/server/src/ws/gateway.ts

import type { FastifyInstance } from 'fastify';
import type { WebSocket } from 'ws';
import {
  BINARY_FRAME,
  PROTOCOL_VERSION,
  encodeBinaryFrame,
  parseClientMessage,
  type RepoSummary,
  type ServerMessage,
  type SessionSummary,
  type WorktreeSummary,
} from '@arboretum/shared';
import type { ClientBinding, PtyManager } from '../core/pty-manager.js';
import type { DiscoveryService } from '../core/discovery-service.js';
import type { WorktreeManager } from '../core/worktree-manager.js';

const HEARTBEAT_MS = 30_000;

interface ChannelState {
  sessionId: string;
  binding: ClientBinding;
}

export function registerWsGateway(
  app: FastifyInstance,
  manager: PtyManager,
  discovery: DiscoveryService,
  worktrees: WorktreeManager,
  serverVersion: string,
): void {
  app.get('/ws', { websocket: true }, (socket: WebSocket, req) => {
    // L'auth + le check Origin ont eu lieu dans le preValidation global (app.ts).
    const channels = new Map<number, ChannelState>();
    let nextChannel = 1;
    let helloDone = false;
    let subscribedSessions = false;
    let subscribedWorktrees = false;
    let alive = true;

    const send = (msg: ServerMessage): void => {
      if (socket.readyState === socket.OPEN) socket.send(JSON.stringify(msg));
    };
    const sendBinary = (type: number, channel: number, payload: Buffer): void => {
      if (socket.readyState === socket.OPEN) socket.send(encodeBinaryFrame(type, channel, payload));
    };

    const onSessionUpdate = (session: SessionSummary): void => {
      if (subscribedSessions) send({ type: 'session_update', session });
    };
    const onSessionExit = (e: { sessionId: string; exitCode: number | null; signal: number | null }): void => {
      if (subscribedSessions) send({ type: 'session_exit', ...e });
    };
    // Les sessions découvertes (DiscoveryService) sont relayées sur le même flux que les managées.
    const onDiscoveryUpdate = (session: SessionSummary): void => {
      if (subscribedSessions) send({ type: 'session_update', session });
    };
    const onRepoUpdate = (repo: RepoSummary): void => {
      if (subscribedWorktrees) send({ type: 'repo_update', repo });
    };
    const onRepoRemoved = (repoId: string): void => {
      if (subscribedWorktrees) send({ type: 'repo_removed', repoId });
    };
    const onWorktreeUpdate = (e: { repoId: string; worktree: WorktreeSummary }): void => {
      if (subscribedWorktrees) send({ type: 'worktree_update', ...e });
    };
    const onWorktreeRemoved = (e: { repoId: string; path: string }): void => {
      if (subscribedWorktrees) send({ type: 'worktree_removed', ...e });
    };
    manager.on('session_update', onSessionUpdate);
    manager.on('session_exit', onSessionExit);
    discovery.on('discovery_update', onDiscoveryUpdate);
    worktrees.on('repo_update', onRepoUpdate);
    worktrees.on('repo_removed', onRepoRemoved);
    worktrees.on('worktree_update', onWorktreeUpdate);
    worktrees.on('worktree_removed', onWorktreeRemoved);

    const heartbeat = setInterval(() => {
      if (!alive) {
        socket.terminate();
        return;
      }
      alive = false;
      socket.ping();
    }, HEARTBEAT_MS);
    socket.on('pong', () => {
      alive = true;
    });

    socket.on('message', (raw, isBinary) => {
      if (isBinary) return; // P1 : aucune frame binaire attendue côté client
      const msg = parseClientMessage(String(raw));
      if (!msg) {
        send({ type: 'error', code: 'BAD_MESSAGE', message: 'Unparseable or invalid message' });
        return;
      }
      if (!helloDone && msg.type !== 'hello') {
        send({ type: 'error', code: 'BAD_PROTOCOL', message: 'Expected hello first' });
        return;
      }
      switch (msg.type) {
        case 'hello': {
          if (msg.protocol !== PROTOCOL_VERSION) {
            send({ type: 'error', code: 'BAD_PROTOCOL', message: `Server speaks protocol ${PROTOCOL_VERSION}` });
            socket.close();
            return;
          }
          helloDone = true;
          send({ type: 'hello_ok', protocol: PROTOCOL_VERSION, serverVersion });
          return;
        }
        case 'sub': {
          subscribedSessions = msg.topics.includes('sessions');
          subscribedWorktrees = msg.topics.includes('worktrees');
          return;
        }
        case 'attach': {
          const channel = nextChannel++;
          const binding: ClientBinding = {
            channel,
            mode: msg.mode,
            controlling: false,
            sentBytes: 0,
            ackedBytes: 0,
            lagging: false,
            sendOutput: (payload) => sendBinary(BINARY_FRAME.OUTPUT, channel, payload),
            sendResync: (payload) => sendBinary(BINARY_FRAME.RESYNC, channel, payload),
            onDetached: (reason) => {
              channels.delete(channel);
              send({ type: 'detached', channel, reason });
            },
            onControlChanged: (controlling) => send({ type: 'control_changed', channel, controlling }),
          };
          const res = manager.attach(msg.sessionId, binding, msg.cols, msg.rows);
          if (!res.ok) {
            send({ type: 'error', code: res.code, message: `Cannot attach: ${res.code}` });
            return;
          }
          channels.set(channel, { sessionId: msg.sessionId, binding });
          send({ type: 'attached', channel, sessionId: msg.sessionId, mode: msg.mode, controlling: res.controlling });
          return;
        }
        case 'detach': {
          const st = channels.get(msg.channel);
          if (!st) return;
          channels.delete(msg.channel);
          manager.detach(st.sessionId, st.binding);
          send({ type: 'detached', channel: msg.channel, reason: 'client' });
          return;
        }
        case 'stdin': {
          const st = channels.get(msg.channel);
          if (!st) {
            send({ type: 'error', code: 'NOT_ATTACHED', message: 'Unknown channel', channel: msg.channel });
            return;
          }
          const r = manager.write(st.sessionId, st.binding, msg.data);
          if (r === 'not_controlling') {
            send({ type: 'error', code: 'NOT_CONTROLLING', message: 'Observer mode is read-only', channel: msg.channel });
          } else if (r === 'gone') {
            send({ type: 'error', code: 'SESSION_EXITED', message: 'Session has exited', channel: msg.channel });
          }
          return;
        }
        case 'resize': {
          const st = channels.get(msg.channel);
          if (st) manager.resize(st.sessionId, st.binding, msg.cols, msg.rows);
          return;
        }
        case 'ack': {
          const st = channels.get(msg.channel);
          if (st) manager.ack(st.sessionId, st.binding, msg.bytes);
          return;
        }
        case 'ping':
          send({ type: 'pong' });
          return;
      }
    });

    socket.on('close', () => {
      clearInterval(heartbeat);
      manager.off('session_update', onSessionUpdate);
      manager.off('session_exit', onSessionExit);
      discovery.off('discovery_update', onDiscoveryUpdate);
      worktrees.off('repo_update', onRepoUpdate);
      worktrees.off('repo_removed', onRepoRemoved);
      worktrees.off('worktree_update', onWorktreeUpdate);
      worktrees.off('worktree_removed', onWorktreeRemoved);
      for (const [, st] of channels) manager.detach(st.sessionId, st.binding);
      channels.clear();
    });

    void req;
  });
}