arboretum/packages/server/src/core/pty-manager.ts

import { EventEmitter } from 'node:events';
import { existsSync, statSync } from 'node:fs';
import { randomUUID } from 'node:crypto';
import pty from '@homebridge/node-pty-prebuilt-multiarch';
import { FLOW, REPLAY_TAIL_BYTES, type SessionSummary } from '@arboretum/shared';
import { RingBuffer } from './ring-buffer.js';
import { buildSpawnSpec } from './claude-launcher.js';
import type { Db } from '../db/index.js';

const RING_CAPACITY = 2 * 1024 * 1024;
const KILL_GRACE_MS = 5000;

/** Lien entre un client WS attaché et une session. La gateway fournit les callbacks d'envoi. */
export interface ClientBinding {
  channel: number;
  mode: 'interactive' | 'observer';
  controlling: boolean;
  sentBytes: number;
  ackedBytes: number;
  lagging: boolean;
  sendOutput: (payload: Buffer) => void;
  sendResync: (payload: Buffer) => void;
  onDetached: (reason: 'session_exit' | 'replaced') => void;
  onControlChanged: (controlling: boolean) => void;
}

interface ManagedSession {
  id: string;
  cwd: string;
  command: 'claude' | 'bash';
  title: string | null;
  createdAt: string;
  proc: pty.IPty;
  ring: RingBuffer;
  clients: Set<ClientBinding>;
  paused: boolean;
  exited: { exitCode: number | null; signal: number | null } | null;
  killTimer: NodeJS.Timeout | null;
}

export interface PtyManagerEvents {
  session_update: [SessionSummary];
  session_exit: [{ sessionId: string; exitCode: number | null; signal: number | null }];
}

export class PtyManager extends EventEmitter<PtyManagerEvents> {
  private readonly live = new Map<string, ManagedSession>();

  constructor(private readonly db: Db) {
    super();
  }

  spawn(opts: { cwd: string; command?: 'claude' | 'bash' }): SessionSummary {
    const cwd = opts.cwd;
    if (!existsSync(cwd) || !statSync(cwd).isDirectory()) {
      throw Object.assign(new Error(`Not a directory: ${cwd}`), { statusCode: 400 });
    }
    const command = opts.command ?? 'claude';
    const spec = buildSpawnSpec(command);
    const id = randomUUID();
    const proc = pty.spawn(spec.file, spec.args, {
      name: 'xterm-256color',
      cols: 120,
      rows: 32,
      cwd,
      env: spec.env as { [key: string]: string },
    });
    const session: ManagedSession = {
      id,
      cwd,
      command,
      title: null,
      createdAt: new Date().toISOString(),
      proc,
      ring: new RingBuffer(RING_CAPACITY),
      clients: new Set(),
      paused: false,
      exited: null,
      killTimer: null,
    };
    this.live.set(id, session);
    this.db
      .prepare('INSERT INTO sessions (id, cwd, command, created_at) VALUES (?, ?, ?, ?)')
      .run(id, cwd, command, session.createdAt);

    proc.onData((data) => this.handleOutput(session, Buffer.from(data, 'utf8')));
    proc.onExit(({ exitCode, signal }) => this.handleExit(session, exitCode, signal ?? null));

    const summary = this.summarize(session);
    this.emit('session_update', summary);
    return summary;
  }

  list(): SessionSummary[] {
    const liveSummaries = [...this.live.values()].map((s) => this.summarize(s));
    const liveIds = new Set(this.live.keys());
    const rows = this.db
      .prepare('SELECT id, cwd, command, title, created_at, ended_at, exit_code FROM sessions ORDER BY created_at DESC LIMIT 100')
      .all() as Array<{ id: string; cwd: string; command: string; title: string | null; created_at: string; ended_at: string | null; exit_code: number | null }>;
    const historical: SessionSummary[] = rows
      .filter((r) => !liveIds.has(r.id))
      .map((r) => ({
        id: r.id,
        cwd: r.cwd,
        command: r.command,
        title: r.title,
        status: 'exited',
        live: false,
        createdAt: r.created_at,
        endedAt: r.ended_at,
        exitCode: r.exit_code,
        clients: 0,
      }));
    return [...liveSummaries, ...historical];
  }

  get(id: string): SessionSummary | null {
    const s = this.live.get(id);
    return s ? this.summarize(s) : null;
  }

  kill(id: string): boolean {
    const s = this.live.get(id);
    if (!s || s.exited) return false;
    try {
      process.kill(s.proc.pid, 'SIGTERM');
    } catch {
      return false;
    }
    s.killTimer ??= setTimeout(() => {
      if (!s.exited) {
        try {
          process.kill(s.proc.pid, 'SIGKILL');
        } catch {
          /* déjà mort */
        }
      }
    }, KILL_GRACE_MS);
    return true;
  }

  /** Arrêt du daemon : SIGTERM à toutes les sessions (le CLI nettoie son registre sur SIGTERM — spike S1). */
  shutdown(): void {
    for (const id of this.live.keys()) this.kill(id);
  }

  // ---- attach / detach / io ----

  attach(sessionId: string, binding: ClientBinding, cols: number, rows: number): { ok: true; controlling: boolean } | { ok: false; code: 'NOT_FOUND' | 'SESSION_EXITED' } {
    const s = this.live.get(sessionId);
    if (!s) return { ok: false, code: 'NOT_FOUND' };
    if (s.exited) return { ok: false, code: 'SESSION_EXITED' };
    const hasController = [...s.clients].some((c) => c.controlling);
    binding.controlling = binding.mode === 'interactive' && !hasController;
    s.clients.add(binding);
    if (binding.controlling) s.proc.resize(cols, rows);
    // Replay : reset terminal + queue du ring (l'écran TUI courant se reconstitue)
    binding.sendResync(s.ring.tail(REPLAY_TAIL_BYTES));
    binding.sentBytes = 0;
    binding.ackedBytes = 0;
    return { ok: true, controlling: binding.controlling };
  }

  detach(sessionId: string, binding: ClientBinding): void {
    const s = this.live.get(sessionId);
    if (!s) return;
    s.clients.delete(binding);
    if (binding.controlling) {
      const next = [...s.clients].find((c) => c.mode === 'interactive');
      if (next) {
        next.controlling = true;
        next.onControlChanged(true);
      }
    }
    this.updateFlowControl(s);
  }

  write(sessionId: string, binding: ClientBinding, data: string): 'ok' | 'not_controlling' | 'gone' {
    const s = this.live.get(sessionId);
    if (!s || s.exited) return 'gone';
    // Outil mono-utilisateur : tous les interactifs peuvent écrire ; les observers jamais.
    if (binding.mode !== 'interactive') return 'not_controlling';
    s.proc.write(data);
    return 'ok';
  }

  resize(sessionId: string, binding: ClientBinding, cols: number, rows: number): void {
    const s = this.live.get(sessionId);
    if (!s || s.exited || !binding.controlling) return;
    s.proc.resize(cols, rows);
  }

  ack(sessionId: string, binding: ClientBinding, bytes: number): void {
    const s = this.live.get(sessionId);
    if (!s) return;
    binding.ackedBytes = Math.max(binding.ackedBytes, bytes);
    if (binding.lagging && binding.sentBytes - binding.ackedBytes < FLOW.LOW_WATERMARK) {
      binding.lagging = false;
      binding.sendResync(s.ring.tail(REPLAY_TAIL_BYTES));
      binding.sentBytes = 0;
      binding.ackedBytes = 0;
    }
    this.updateFlowControl(s);
  }

  // ---- interne ----

  private handleOutput(s: ManagedSession, chunk: Buffer): void {
    s.ring.write(chunk);
    for (const c of s.clients) {
      if (c.lagging) continue;
      c.sendOutput(chunk);
      c.sentBytes += chunk.length;
      if (c.sentBytes - c.ackedBytes > FLOW.LAGGING_BYTES) c.lagging = true;
    }
    this.updateFlowControl(s);
  }

  /**
   * pause() seulement quand TOUS les clients interactifs non-lagging dépassent HIGH ;
   * resume() quand le min repasse sous LOW. Les observers ne freinent jamais le PTY.
   */
  private updateFlowControl(s: ManagedSession): void {
    if (s.exited) return;
    const interactive = [...s.clients].filter((c) => c.mode === 'interactive' && !c.lagging);
    if (interactive.length === 0) {
      if (s.paused) {
        s.proc.resume();
        s.paused = false;
      }
      return;
    }
    const outstandings = interactive.map((c) => c.sentBytes - c.ackedBytes);
    const min = Math.min(...outstandings);
    if (!s.paused && outstandings.every((o) => o > FLOW.HIGH_WATERMARK)) {
      s.proc.pause();
      s.paused = true;
    } else if (s.paused && min < FLOW.LOW_WATERMARK) {
      s.proc.resume();
      s.paused = false;
    }
  }

  private handleExit(s: ManagedSession, exitCode: number | null, signal: number | null): void {
    s.exited = { exitCode, signal };
    if (s.killTimer) clearTimeout(s.killTimer);
    this.db
      .prepare('UPDATE sessions SET ended_at = ?, exit_code = ? WHERE id = ?')
      .run(new Date().toISOString(), exitCode, s.id);
    for (const c of s.clients) c.onDetached('session_exit');
    s.clients.clear();
    this.live.delete(s.id);
    this.emit('session_exit', { sessionId: s.id, exitCode, signal });
    this.emit('session_update', { ...this.summarize(s), live: false, status: 'exited' });
  }

  private summarize(s: ManagedSession): SessionSummary {
    return {
      id: s.id,
      cwd: s.cwd,
      command: s.command,
      title: s.title,
      status: s.exited ? 'exited' : 'running',
      live: !s.exited,
      createdAt: s.createdAt,
      endedAt: null,
      exitCode: s.exited?.exitCode ?? null,
      clients: s.clients.size,
    };
  }
}