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

import type { Server } from 'node:http';
import type { Duplex } from 'node:stream';
import type { WebSocket } from 'ws';
import { WebSocketServer } from 'ws';

import type { BoxHandle } from './user-container';
import { boxHomePaths } from './box';
import { env } from './env';
import { getDockerClient } from './runtime/docker';
import {
  parseBoxTokenUsername,
  verifyBoxTerminalToken,
  verifyTerminalToken,
} from './terminal-token';
import { acquireUserBox } from './user-container';
import { ensureBashProfile } from './user-environment';
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);
};

export const parseResizeMessage = (
  data: Buffer,
  isBinary: boolean,
): { cols: number; rows: number } | null => {
  if (isBinary) return null;
  try {
    const message = JSON.parse(data.toString('utf8')) as {
      type?: string;
      cols?: number;
      rows?: number;
    };
    if (message.type !== 'resize') return null;
    const cols = clampDimension(message.cols);
    const rows = clampDimension(message.rows);
    if (!cols || !rows) return null;
    return { cols, rows };
  } catch {
    return null;
  }
};

// The login-shell command shared by the job and box terminals: prefer a
// resumable tmux session, else fall back to a plain interactive login shell.
const SHELL_CMD = [
  '/bin/bash',
  '-lc',
  'if command -v tmux >/dev/null 2>&1; then exec tmux new-session -A -s spoon; else exec bash -il; fi',
];

type ShellBridgeOptions = {
  // Acquire the per-user box (reference-counted); its handle is released on
  // cleanup. Called after the input handler is wired so early keystrokes buffer.
  acquire: () => Promise<BoxHandle>;
  // Host-side preparation before the shell starts (e.g. seed .bash_profile).
  prepare?: () => Promise<void>;
  // Directory the login shell starts in (container path).
  cwd: string;
  // Env entries appended after TERM (e.g. HOME, and job secrets for job terms).
  envFlags: string[];
};

/**
 * Shared PTY bridge: wire a WebSocket to an interactive login shell running via
 * a real TTY (dockerode `exec` with `Tty:true` + `exec.resize`) inside the
 * user's box. Registers the input handler immediately and buffers 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. Used by both the
 * job terminal (`bridge`) and the box terminal (`bridgeBox`); only the box to
 * acquire, the working directory, and the env differ.
 */
const runShellBridge = async (ws: WebSocket, opts: ShellBridgeOptions) => {
  const execHolder: {
    current?: { exec: import('dockerode').Exec; stream: Duplex };
  } = {};
  const pendingInput: Buffer[] = [];
  let cols = 80;
  let rows = 24;

  ws.on('message', (data: Buffer, isBinary: boolean) => {
    const resize = parseResizeMessage(data, isBinary);
    if (resize) {
      cols = resize.cols;
      rows = resize.rows;
      void execHolder.current?.exec
        .resize({ h: rows, w: cols })
        .catch(() => undefined);
      return;
    }
    if (execHolder.current) execHolder.current.stream.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;
    execHolder.current?.stream.end();
    execHolder.current?.stream.destroy();
    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 {
    if (opts.prepare) await opts.prepare();
    const handle = await opts.acquire();
    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;
  }

  const docker = getDockerClient();
  const container = docker.getContainer(boxName);
  const exec = await container.exec({
    AttachStdin: true,
    AttachStdout: true,
    AttachStderr: true,
    Tty: true,
    Cmd: SHELL_CMD,
    Env: ['TERM=xterm-256color', ...opts.envFlags],
    WorkingDir: opts.cwd,
  });
  const stream = await exec.start({ hijack: true, stdin: true, Tty: true });
  execHolder.current = { exec, stream };

  if (isReleased()) {
    stream.end();
    stream.destroy();
    return;
  }

  await exec.resize({ h: rows, w: cols }).catch(() => undefined);

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

  stream.on('data', (chunk: Buffer) => {
    if (ws.readyState === ws.OPEN) ws.send(chunk, { binary: true });
  });
  stream.on('end', () => {
    if (ws.readyState === ws.OPEN) ws.close();
  });
  stream.on('error', () => {
    if (ws.readyState === ws.OPEN) ws.close();
  });
};

// Job terminal: opens the shell at the repo checkout, with the job's selected
// secrets in the environment (it is scoped to that job's active workspace).
const bridge = async (ws: WebSocket, jobId: string) => {
  const workspace = getTerminalWorkspace(jobId);
  if (!workspace) {
    ws.close(1011, 'Workspace is not active.');
    return;
  }
  await runShellBridge(ws, {
    acquire: () =>
      acquireUserBox({
        username: workspace.username,
        workdir: workspace.workdir,
        containerHome: workspace.containerHome,
      }),
    cwd: workspace.containerRepo,
    envFlags: [
      `HOME=${workspace.containerHome}`,
      ...workspace.secrets.map((secret) => `${secret.name}=${secret.value}`),
    ],
  });
};

// Box terminal: opens a login shell rooted at the user's home (`~`), with no
// per-job secrets — it is user-scoped, not tied to any job's workspace.
const bridgeBox = async (ws: WebSocket, username: string) => {
  const { homeDir, containerHome } = boxHomePaths(username);
  await runShellBridge(ws, {
    prepare: () => ensureBashProfile(homeDir),
    acquire: () =>
      acquireUserBox({ username, workdir: homeDir, containerHome }),
    cwd: containerHome,
    envFlags: [`HOME=${containerHome}`],
  });
};

/**
 * Attaches the interactive-terminal WebSocket endpoints to the worker's HTTP
 * server. Two routes, each authorized by a short-lived token minted by the Next
 * app: `/jobs/:jobId/terminal?token=…` (job-scoped, after ownership check) and
 * `/box/terminal?token=…` (user-scoped, username read from the token itself).
 */
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 token = url.searchParams.get('token') ?? '';

    // User-scoped box terminal: the token authorizes the connection, so the
    // username is read FROM the token, then verified against the whole token.
    if (url.pathname === '/box/terminal') {
      const username = parseBoxTokenUsername(token);
      if (
        !username ||
        !verifyBoxTerminalToken(token, username, env.terminalSecret)
      ) {
        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
        socket.destroy();
        return;
      }
      wss.handleUpgrade(request, socket, head, (ws) => {
        void bridgeBox(ws, username).catch(() => {
          if (ws.readyState === ws.OPEN) ws.close();
        });
      });
      return;
    }

    const match = /^\/jobs\/([^/]+)\/terminal$/.exec(url.pathname);
    if (!match?.[1]) {
      socket.destroy();
      return;
    }
    const jobId = decodeURIComponent(match[1]);
    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).catch(() => {
        if (ws.readyState === ws.OPEN) ws.close();
      });
    });
  });

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