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 GroupSummary,
  type RepoSummary,
  type CloneOperation,
  type ServerMessage,
  type SessionSummary,
  type SettingsBroadcast,
  type WorktreeSummary,
} from '@arboretum/shared';
import type { ClientBinding, PtyManager } from '../core/pty-manager.js';
import type { DiscoveryService } from '../core/discovery-service.js';
import type { SessionArchiveService } from '../core/session-archive.js';
import type { WorktreeManager } from '../core/worktree-manager.js';
import type { GroupManager } from '../core/group-manager.js';
import type { SettingsBus } from '../core/settings-bus.js';
import type { CloneManager } from '../core/clone-manager.js';

const HEARTBEAT_MS = 30_000;

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

export function registerWsGateway(
  app: FastifyInstance,
  manager: PtyManager,
  discovery: DiscoveryService,
  sessionArchive: SessionArchiveService,
  worktrees: WorktreeManager,
  groups: GroupManager,
  settingsBus: SettingsBus,
  clones: CloneManager,
  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 subscribedGroups = false;
    let subscribedSettings = false;
    let subscribedClones = false;
    let alive = true;
    // Worktrees « regardés » par CETTE connexion (clé repoId\0path) → push ciblé de worktree_changes.
    const watched = new Set<string>();
    const watchKey = (repoId: string, path: string): string => `${repoId}\0${path}`;

    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 });
    };
    // P10 · une session managée vient d'être archivée (auto par ancienneté ou manuellement).
    const onSessionArchived = (e: { id: string; archivedAt: string }): void => {
      if (subscribedSessions) send({ type: 'session_archived', sessionId: e.id });
    };
    // 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 });
    };
    const onGroupUpdate = (group: GroupSummary): void => {
      if (subscribedGroups) send({ type: 'group_update', group });
    };
    const onGroupRemoved = (groupId: string): void => {
      if (subscribedGroups) send({ type: 'group_removed', groupId });
    };
    // P11 · un réglage a changé : relayé aux connexions abonnées au topic 'settings'.
    const onSettingsUpdate = (settings: SettingsBroadcast): void => {
      if (subscribedSettings) send({ type: 'settings_update', settings });
    };
    // P12 · progression d'un clone : relayée aux abonnés du topic 'clones'.
    const onCloneUpdate = (operation: CloneOperation): void => {
      if (subscribedClones) send({ type: 'clone_update', operation });
    };
    // P7 · détail d'un worktree modifié : poussé UNIQUEMENT si cette connexion le regarde.
    const onWorktreeChanges = (e: { repoId: string; path: string }): void => {
      if (watched.has(watchKey(e.repoId, e.path))) send({ type: 'worktree_changes', ...e });
    };
    manager.on('session_update', onSessionUpdate);
    manager.on('session_exit', onSessionExit);
    sessionArchive.on('session_archived', onSessionArchived);
    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);
    worktrees.on('worktree_changes', onWorktreeChanges);
    groups.on('group_update', onGroupUpdate);
    groups.on('group_removed', onGroupRemoved);
    settingsBus.on('settings_update', onSettingsUpdate);
    clones.on('clone_update', onCloneUpdate);

    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');
          subscribedGroups = msg.topics.includes('groups');
          subscribedSettings = msg.topics.includes('settings');
          subscribedClones = msg.topics.includes('clones');
          return;
        }
        case 'watch': {
          const key = watchKey(msg.repoId, msg.path);
          if (watched.has(key)) return;
          watched.add(key);
          // Validation repo/worktree + armement du watcher FS (asynchrone, best-effort).
          void worktrees.watch(msg.repoId, msg.path).catch((err) => {
            watched.delete(key);
            send({ type: 'error', code: 'NOT_FOUND', message: (err as Error).message });
          });
          return;
        }
        case 'unwatch': {
          watched.delete(watchKey(msg.repoId, msg.path));
          worktrees.unwatch(msg.repoId, msg.path);
          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 'answer': {
          const st = channels.get(msg.channel);
          if (!st) {
            send({ type: 'error', code: 'NOT_ATTACHED', message: 'Unknown channel', channel: msg.channel });
            return;
          }
          const r = manager.answer(st.sessionId, st.binding, msg.action, msg.optionN);
          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 });
          } else if (r === 'invalid') {
            send({ type: 'error', code: 'INVALID_ANSWER', message: 'No such option in the current dialog', 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);
      sessionArchive.off('session_archived', onSessionArchived);
      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);
      worktrees.off('worktree_changes', onWorktreeChanges);
      groups.off('group_update', onGroupUpdate);
      groups.off('group_removed', onGroupRemoved);
      settingsBus.off('settings_update', onSettingsUpdate);
      clones.off('clone_update', onCloneUpdate);
      for (const [, st] of channels) manager.detach(st.sessionId, st.binding);
      channels.clear();
      // Libère le refcount du watcher pour chaque worktree regardé par cette connexion.
      for (const key of watched) {
        const sep = key.indexOf('\0');
        worktrees.unwatch(key.slice(0, sep), key.slice(sep + 1));
      }
      watched.clear();
    });

    void req;
  });
}