spoon/apps/agent-worker/src/terminal.ts

import { spawn } from 'node:child_process';
import type { ChildProcessWithoutNullStreams } from 'node:child_process';
import type { Server } from 'node:http';
import type { WebSocket } from 'ws';
import { WebSocketServer } from 'ws';

import type { BoxHandle } from './user-container';
import { env } from './env';
import { verifyTerminalToken } from './terminal-token';
import { acquireUserBox } from './user-container';
import { getTerminalWorkspace } from './worker';

const clampDimension = (value: unknown) => {
  const n = Math.trunc(Number(value));
  if (!Number.isFinite(n)) return undefined;
  return Math.min(Math.max(n, 1), 1000);
};

// Single-quote a string for a POSIX shell.
const shellQuote = (value: string) => `'${value.replaceAll("'", `'\\''`)}'`;

const bridge = async (ws: WebSocket, jobId: string) => {
  const workspace = getTerminalWorkspace(jobId);
  if (!workspace) {
    ws.close(1011, 'Workspace is not active.');
    return;
  }

  // bun can't load node-pty (native ABI mismatch) and dockerode can't attach to
  // podman, so we drive the runtime CLI (`<runtime> exec -i`) and allocate the PTY
  // *inside* the container with `script`, bridging the plain pipes to the socket.
  //
  // Register the message handler immediately and buffer input/size until the exec
  // is ready (acquiring the box can take seconds on first connect), so the initial
  // resize and early keystrokes aren't dropped.
  const procHolder: { current?: ChildProcessWithoutNullStreams } = {};
  const pendingInput: Buffer[] = [];
  let cols = 80;
  let rows = 24;

  ws.on('message', (data: Buffer, isBinary: boolean) => {
    if (!isBinary) {
      // Text frames are control messages (resize); anything else is raw input.
      try {
        const message = JSON.parse(data.toString('utf8')) as {
          type?: string;
          cols?: number;
          rows?: number;
        };
        if (message.type === 'resize') {
          const c = clampDimension(message.cols);
          const r = clampDimension(message.rows);
          if (c && r) {
            cols = c;
            rows = r;
          }
          return;
        }
      } catch {
        // fall through: treat as raw input
      }
    }
    if (procHolder.current) procHolder.current.stdin.write(data);
    else pendingInput.push(data);
  });

  const handleHolder: { current?: BoxHandle } = {};
  let released = false;
  // Read through a function so TS doesn't narrow `released` to a constant — the
  // cleanup handler flips it asynchronously when the socket closes.
  const isReleased = () => released;
  const cleanup = () => {
    if (released) return;
    released = true;
    procHolder.current?.kill();
    handleHolder.current?.release();
  };
  ws.on('close', cleanup);
  ws.on('error', cleanup);

  // Hold the per-user box open while this terminal is connected; the agent and
  // the terminal share the exact same container (Phase 2).
  let boxName: string;
  try {
    const handle = await acquireUserBox({
      username: workspace.username,
      workdir: workspace.workdir,
      containerHome: workspace.containerHome,
    });
    handleHolder.current = handle;
    boxName = handle.boxName;
  } catch (error) {
    ws.close(
      1011,
      `Failed to start terminal: ${error instanceof Error ? error.message : 'unknown error'}`,
    );
    return;
  }

  // Cleanup may have run before the awaited handle existed.
  if (isReleased()) {
    handleHolder.current.release();
    return;
  }

  // Reattach a persistent tmux session across reconnects when available, else a
  // plain login shell. `stty` sizes the PTY to the client's viewport up front.
  const launcher =
    `stty rows ${rows} cols ${cols} 2>/dev/null; ` +
    // Reattach a persistent tmux session when tmux is present; otherwise fall back
    // to an interactive login shell (`-i` so it prints a prompt and line-edits).
    // Check with `command -v` rather than `exec tmux || …`: a failed `exec` makes a
    // non-interactive shell exit before the `||`, so the fallback never runs.
    'if command -v tmux >/dev/null 2>&1; then exec tmux new-session -A -s spoon; ' +
    'else exec bash -il; fi';
  const envFlags = [
    '-e',
    'TERM=xterm-256color',
    '-e',
    `HOME=${workspace.containerHome}`,
    ...workspace.secrets.flatMap((s) => ['-e', `${s.name}=${s.value}`]),
  ];

  const proc = spawn(
    env.containerRuntime,
    [
      'exec',
      '-i',
      ...envFlags,
      '-w',
      workspace.containerRepo,
      boxName,
      '/bin/bash',
      '-lc',
      `exec script -qfc ${shellQuote(launcher)} /dev/null`,
    ],
    { stdio: ['pipe', 'pipe', 'pipe'] },
  );
  procHolder.current = proc;

  // Replay any keystrokes the client sent before the process was ready.
  for (const buffered of pendingInput) proc.stdin.write(buffered);
  pendingInput.length = 0;

  const forward = (chunk: Buffer) => {
    if (ws.readyState === ws.OPEN) ws.send(chunk, { binary: true });
  };
  proc.stdout.on('data', forward);
  proc.stderr.on('data', forward);
  proc.on('exit', () => {
    if (ws.readyState === ws.OPEN) ws.close();
  });
  proc.on('error', () => {
    if (ws.readyState === ws.OPEN) ws.close();
  });
};

/**
 * Attaches the interactive-terminal WebSocket endpoint to the worker's HTTP
 * server. Browser connects to `/jobs/:jobId/terminal?token=…` with a short-lived
 * token minted by the Next app (which has already verified job ownership).
 */
export const attachTerminalServer = (server: Server) => {
  if (env.runtime !== 'docker') return;
  const wss = new WebSocketServer({ noServer: true });

  server.on('upgrade', (request, socket, head) => {
    const url = new URL(request.url ?? '', `http://localhost:${env.httpPort}`);
    const match = /^\/jobs\/([^/]+)\/terminal$/.exec(url.pathname);
    if (!match?.[1]) {
      socket.destroy();
      return;
    }
    const jobId = decodeURIComponent(match[1]);
    const token = url.searchParams.get('token') ?? '';
    if (!verifyTerminalToken(token, jobId, env.terminalSecret)) {
      socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
      socket.destroy();
      return;
    }
    wss.handleUpgrade(request, socket, head, (ws) => {
      void bridge(ws, jobId);
    });
  });

  console.log('Spoon agent worker terminal WebSocket endpoint enabled.');
};