homelab/apps/harness/pty-server.js

// @ts-check
const pty = require("node-pty");
const { WebSocketServer } = require("ws");

const MAX_SESSIONS = 3;
const MAX_SESSION_AGE_MS = 2 * 60 * 60 * 1000; // 2 hours
const CLEANUP_INTERVAL_MS = 5 * 60 * 1000; // 5 minutes

/** @type {Map<string, { pty: import("node-pty").IPty, ws: import("ws").WebSocket, createdAt: number }>} */
const sessions = new Map();

/**
 * @param {import("http").Server} server
 */
function attachPtyWebSocket(server) {
  const wss = new WebSocketServer({ noServer: true });

  server.on("upgrade", (req, socket, head) => {
    const url = new URL(req.url || "/", `http://${req.headers.host}`);
    if (url.pathname !== "/ws/pty") return;

    wss.handleUpgrade(req, socket, head, (ws) => {
      wss.emit("connection", ws, req);
    });
  });

  wss.on("connection", async (ws, req) => {
    const url = new URL(req.url || "/", `http://${req.headers.host}`);
    const agentId = url.searchParams.get("agentId");
    const cols = parseInt(url.searchParams.get("cols") || "80", 10);
    const rows = parseInt(url.searchParams.get("rows") || "24", 10);

    if (!agentId) {
      ws.send(JSON.stringify({ type: "error", message: "agentId required" }));
      ws.close();
      return;
    }

    if (sessions.size >= MAX_SESSIONS) {
      ws.send(
        JSON.stringify({
          type: "error",
          message: `max ${MAX_SESSIONS} concurrent sessions`,
        }),
      );
      ws.close();
      return;
    }

    // Fetch agent config + env from the Next.js API
    const port = process.env.PORT || 3100;
    let config;
    try {
      const res = await fetch(
        `http://127.0.0.1:${port}/api/agents/${encodeURIComponent(agentId)}/env`,
      );
      if (!res.ok) {
        const body = await res.json().catch(() => ({}));
        ws.send(
          JSON.stringify({
            type: "error",
            message: body.error || `agent lookup failed: ${res.status}`,
          }),
        );
        ws.close();
        return;
      }
      config = await res.json();
    } catch (err) {
      ws.send(
        JSON.stringify({
          type: "error",
          message: `failed to fetch agent config: ${err.message}`,
        }),
      );
      ws.close();
      return;
    }

    const sessionId = `${agentId}-${Date.now()}`;
    let term;
    try {
      term = pty.spawn(config.command, config.args, {
        name: "xterm-256color",
        cols,
        rows,
        cwd: config.workDir,
        env: config.env,
      });
    } catch (err) {
      ws.send(
        JSON.stringify({
          type: "error",
          message: `failed to spawn PTY: ${err.message}`,
        }),
      );
      ws.close();
      return;
    }

    sessions.set(sessionId, { pty: term, ws, createdAt: Date.now() });

    ws.send(JSON.stringify({ type: "connected", sessionId }));

    term.onData((data) => {
      if (ws.readyState === ws.OPEN) {
        ws.send(data);
      }
    });

    term.onExit(({ exitCode }) => {
      if (ws.readyState === ws.OPEN) {
        ws.send(JSON.stringify({ type: "exit", code: exitCode }));
        ws.close();
      }
      sessions.delete(sessionId);
    });

    ws.on("message", (data) => {
      const msg = data.toString();
      // Try to parse as JSON control message
      if (msg.startsWith("{")) {
        try {
          const ctrl = JSON.parse(msg);
          if (ctrl.type === "resize" && ctrl.cols && ctrl.rows) {
            term.resize(ctrl.cols, ctrl.rows);
            return;
          }
        } catch {
          // Not JSON, treat as terminal input
        }
      }
      term.write(msg);
    });

    ws.on("close", () => {
      try {
        term.kill();
      } catch {
        // Already dead
      }
      sessions.delete(sessionId);
    });
  });

  // Periodic cleanup of stale sessions
  setInterval(() => {
    const now = Date.now();
    for (const [id, session] of sessions) {
      if (now - session.createdAt > MAX_SESSION_AGE_MS) {
        try {
          session.pty.kill();
        } catch {
          // ignore
        }
        try {
          session.ws.close();
        } catch {
          // ignore
        }
        sessions.delete(id);
      }
    }
  }, CLEANUP_INTERVAL_MS);

  return wss;
}

module.exports = { attachPtyWebSocket };