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