boocode/apps/booterm/src/ws/attach.ts

import type { FastifyInstance } from 'fastify';
import type { IPty } from 'node-pty';
import { getSessionInfo } from '../db.js';
import {
  sanitizeId,
  tmuxSessionName,
  ensureSession,
  capturePane,
} from '../pty/manager.js';
import { attachPty } from '../pty/pty.js';
import { getUser } from '../auth.js';
import { register, unregister } from '../pty/registry.js';

export function registerWsAttachRoute(app: FastifyInstance, tmuxConfPath: string): void {
  app.get<{
    Params: { sid: string; pid: string };
    Querystring: { cols?: string; rows?: string };
  }>(
    '/ws/term/sessions/:sid/panes/:pid',
    { websocket: true },
    async (socket, req) => {
      const sid = sanitizeId(req.params.sid);
      const pid = sanitizeId(req.params.pid);
      if (!sid || !pid) {
        socket.close(1008, 'bad_id_format');
        return;
      }

      const user = getUser(req);
      req.log.info({ user, sid, pid }, 'ws attach');

      const session = await getSessionInfo(sid);
      if (!session) {
        socket.close(1008, 'unknown_session');
        return;
      }

      const sessionName = tmuxSessionName(pid);
      const cols = parseInt(req.query.cols ?? '', 10) || 80;
      const rows = parseInt(req.query.rows ?? '', 10) || 24;

      // Idempotent — /start typically created the session already, but cover
      // the race where the client opens the WS before /start's response lands
      // (or skips /start entirely). With per-pane tmux sessions there's no
      // cross-pane interference, so creating-on-attach is safe.
      try {
        await ensureSession(
          tmuxConfPath,
          sessionName,
          session.project_path,
          req.log,
          cols,
          rows,
        );
      } catch (err) {
        req.log.error({ err }, 'ensureSession failed in WS handler');
        socket.close(1011, 'tmux_failed');
        return;
      }

      register(sid, pid, session.project_path);

      let handle: IPty;
      try {
        handle = attachPty({
          sessionName,
          projectRoot: session.project_path,
          cols,
          rows,
          tmuxConfPath,
        });
      } catch (err) {
        req.log.error({ err }, 'attachPty failed');
        socket.close(1011, 'pty_spawn_failed');
        return;
      }

      // Frame contract (boolab pattern):
      //   server → client text:    JSON control — `init` on connect, `exit` on PTY death
      //   server → client binary:  raw PTY bytes (first frame after init = capture-pane replay)
      //   client → server binary:  user keystrokes
      //   client → server text:    JSON control — `{type:"resize", cols, rows}`
      //
      // The init frame lets the client term.clear() before paint so a remount
      // doesn't show stale buffer content. The capture-pane replay then
      // paints the current tmux pane state into the fresh xterm.
      try {
        socket.send(JSON.stringify({ type: 'init', cols, rows, tmux_session: sessionName }));
      } catch (err) {
        req.log.warn({ err }, 'init frame send failed');
      }

      try {
        const capture = await capturePane(tmuxConfPath, sessionName);
        if (capture.length > 0) {
          socket.send(Buffer.from(capture, 'utf8'), { binary: true });
        }
      } catch (err) {
        req.log.warn({ err }, 'capture-pane failed');
      }

      const onData = (data: string): void => {
        if (socket.readyState !== socket.OPEN) return;
        try {
          socket.send(Buffer.from(data, 'utf8'), { binary: true });
        } catch (err) {
          req.log.warn({ err }, 'ws send failed');
        }
      };
      handle.onData(onData);

      socket.on('message', (rawData: Buffer | string, isBinary?: boolean) => {
        // ws v8 emits Buffer + isBinary boolean; older versions emit string
        // for text frames. Either way: text path tries JSON parse for the
        // resize control; binary path writes to the PTY.
        const isTextFrame = typeof rawData === 'string' || isBinary === false;
        if (isTextFrame) {
          const text = typeof rawData === 'string' ? rawData : rawData.toString('utf8');
          try {
            const parsed = JSON.parse(text) as { type?: string; cols?: number; rows?: number };
            if (parsed.type === 'resize') {
              const newCols = Math.max(1, Math.min(2000, Math.floor(Number(parsed.cols) || 80)));
              const newRows = Math.max(1, Math.min(2000, Math.floor(Number(parsed.rows) || 24)));
              req.log.info({ pid, cols: newCols, rows: newRows }, 'resize');
              try {
                handle.resize(newCols, newRows);
              } catch {
                /* ignore — invalid winsize bubble */
              }
            }
          } catch {
            /* malformed text frame — drop silently */
          }
          return;
        }
        try {
          handle.write((rawData as Buffer).toString('utf8'));
        } catch (err) {
          req.log.warn({ err }, 'pty write failed');
        }
      });

      handle.onExit(({ exitCode }) => {
        try {
          if (socket.readyState === socket.OPEN) {
            socket.send(JSON.stringify({ type: 'exit', code: exitCode }));
          }
        } catch {
          /* ignore */
        }
        try {
          socket.close(1000);
        } catch {
          /* ignore */
        }
      });

      // WS close kills the tmux client (the local PTY) but the tmux server +
      // session persist — so a refresh resumes with full scrollback. Permanent
      // teardown happens via the /kill route called from the frontend when the
      // user closes the pane.
      socket.on('close', () => {
        unregister(pid);
        try {
          handle.kill();
        } catch {
          /* ignore */
        }
      });
    },
  );
}