agentic-dev/apps/runner/src/server.ts

import { once } from "node:events";
import { timingSafeEqual } from "node:crypto";
import { WebSocketServer, type RawData, type WebSocket } from "ws";
import {
  envelopeSchema,
  PROTOCOL_VERSION,
  type RunnerErrorCause,
  type RunnerMessage,
  type RunnerStage,
} from "@repo/core-runner-protocol";
import { runCommand } from "./exec";
import type { Logger } from "./log";
import { cloneRepository } from "./stages/clone";
import { installDependencies } from "./stages/install";
import { runStage, type StageEmitter } from "./stages/stage-runner";

/**
 * The runner's WS protocol server (ADR-027, protocol version "0").
 *
 * Session semantics:
 *
 * - Every inbound frame is gated: JSON → protocol version → envelope
 *   schema (`.strict()`) → workspace-token match. Rejections emit a named
 *   `error` event (`unsupported-protocol-version` / `invalid-message` /
 *   `unauthorized`); pre-handshake rejections and token mismatches also
 *   close the socket (1008 policy violation).
 * - The first accepted message MUST be `hello`; the runner answers
 *   `ready` ("the runner will take commands").
 * - Commands run strictly one at a time in arrival order. While a stage
 *   runs, the runner streams `status` events (stage + elapsedMs). A
 *   command that succeeds ends with `ready` — the idle marker the
 *   orchestrator (story 07) drives the session on. A command that fails
 *   ends with a named `error` event instead; the session stays open so
 *   the client can correct and retry.
 * - Every outbound message is envelope-wrapped and zod-parsed before it
 *   is sent. Outbound envelopes carry the workspace token only once the
 *   peer has proven possession of it; rejection replies to unproven
 *   peers carry a redacted placeholder so the token cannot leak.
 * - Error messages summarize schema issues by path + code only — raw
 *   payload contents (which may include a PAT) are never echoed into
 *   errors or logs.
 */
export interface RunnerServerOptions {
  host: string;
  port: number;
  token: string;
  workspaceDir: string;
  heartbeatMs: number;
  log: Logger;
}

export interface RunnerServer {
  host: string;
  /** The actually-bound port (options.port `0` binds an ephemeral one). */
  port: number;
  /** Terminate all sessions and stop listening. Idempotent. */
  close: () => Promise<void>;
}

/** Placeholder token on replies to peers that never proved possession. */
const REDACTED_TOKEN = "redacted";

function tokensMatch(expected: string, provided: string): boolean {
  const expectedBuffer = Buffer.from(expected);
  const providedBuffer = Buffer.from(provided);
  if (expectedBuffer.length !== providedBuffer.length) return false;
  return timingSafeEqual(expectedBuffer, providedBuffer);
}

type GateRejection = {
  ok: false;
  cause: Extract<
    RunnerErrorCause,
    "unsupported-protocol-version" | "invalid-message" | "unauthorized"
  >;
  detail: string;
};

type GateResult = { ok: true; message: RunnerMessage } | GateRejection;

/** Issue summary that never echoes payload values (a payload can hold a PAT). */
function summarizeIssues(error: import("zod").ZodError): string {
  return error.issues
    .map((issue) => `${issue.path.join(".") || "(root)"}: ${issue.code}`)
    .join("; ");
}

/**
 * Gate one inbound frame: JSON → version → envelope schema → token.
 * Exported for direct unit coverage of every rejection branch.
 */
export function gateInbound(raw: string, expectedToken: string): GateResult {
  let parsed: unknown;
  try {
    parsed = JSON.parse(raw);
  } catch {
    return {
      ok: false,
      cause: "invalid-message",
      detail: "frame is not valid JSON",
    };
  }
  if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
    return {
      ok: false,
      cause: "invalid-message",
      detail: "frame is not a JSON object",
    };
  }
  const version = (parsed as Record<string, unknown>).protocolVersion;
  if (version !== PROTOCOL_VERSION) {
    return {
      ok: false,
      cause: "unsupported-protocol-version",
      detail: `this runner speaks protocol version "${PROTOCOL_VERSION}"`,
    };
  }
  const result = envelopeSchema.safeParse(parsed);
  if (!result.success) {
    return {
      ok: false,
      cause: "invalid-message",
      detail: summarizeIssues(result.error),
    };
  }
  if (!tokensMatch(expectedToken, result.data.token)) {
    return {
      ok: false,
      cause: "unauthorized",
      detail: "workspace token mismatch",
    };
  }
  return { ok: true, message: result.data.message };
}

/** Runner-emitted event types a client must never send. */
const RUNNER_EMITTED_TYPES = new Set(["status", "error", "ready"]);

export type CommandMessage = Extract<
  RunnerMessage,
  { type: "clone" | "install" | "scan" | "adapter-start" | "render-frame" }
>;

/** Stages that later walking-skeleton stories wire up (05+). */
const NOT_IMPLEMENTED: Record<
  Exclude<CommandMessage["type"], "clone" | "install">,
  { cause: RunnerErrorCause; stage?: RunnerStage }
> = {
  scan: { cause: "scan-failed", stage: "scanning" },
  "adapter-start": { cause: "adapter-start-failed", stage: "starting-preview" },
  "render-frame": { cause: "render-failed" },
};

interface CommandContext {
  emit: StageEmitter;
  /** Signal successful command completion — the session's idle marker. */
  ready: () => void;
  options: RunnerServerOptions;
}

async function handleCommand(
  command: CommandMessage,
  ctx: CommandContext,
): Promise<void> {
  const { emit, options } = ctx;
  const stageOptions = {
    emit,
    heartbeatMs: options.heartbeatMs,
    log: options.log,
  };
  switch (command.type) {
    case "clone": {
      const result = await runStage(
        "cloning",
        "clone-failed",
        stageOptions,
        () =>
          cloneRepository({
            workspaceDir: options.workspaceDir,
            runCommand,
            log: options.log,
          })(command),
      );
      if (result.ok) ctx.ready();
      return;
    }
    case "install": {
      const result = await runStage(
        "installing",
        "install-failed",
        stageOptions,
        installDependencies({
          workspaceDir: options.workspaceDir,
          runCommand,
          log: options.log,
        }),
      );
      if (result.ok) ctx.ready();
      return;
    }
    default: {
      const notImplemented = NOT_IMPLEMENTED[command.type];
      emit.error(
        notImplemented.cause,
        `the "${command.type}" command is not implemented by this runner build yet`,
        notImplemented.stage,
      );
    }
  }
}

function rawToString(data: RawData): string {
  if (Array.isArray(data)) return Buffer.concat(data).toString("utf8");
  if (data instanceof ArrayBuffer) return Buffer.from(data).toString("utf8");
  return data.toString("utf8");
}

export async function startRunnerServer(
  options: RunnerServerOptions,
): Promise<RunnerServer> {
  const { host, port, token, log } = options;
  const wss = new WebSocketServer({ host, port });
  await once(wss, "listening");
  const address = wss.address();
  if (address === null || typeof address === "string") {
    wss.close();
    throw new Error("runner server: could not determine bound port");
  }

  wss.on("connection", (ws: WebSocket) => {
    let authenticated = false;
    let commandChain: Promise<void> = Promise.resolve();
    log.info("connection-open", {});

    const send = (message: RunnerMessage): void => {
      // Outbound conformance gate: nothing leaves unparsed.
      const envelope = envelopeSchema.parse({
        protocolVersion: PROTOCOL_VERSION,
        token: authenticated ? token : REDACTED_TOKEN,
        message,
      });
      ws.send(JSON.stringify(envelope));
    };

    const emit: StageEmitter = {
      status: (stage, elapsedMs) => send({ type: "status", stage, elapsedMs }),
      error: (cause, message, stage) =>
        send(
          stage
            ? { type: "error", cause, message, stage }
            : { type: "error", cause, message },
        ),
    };

    ws.on("message", (data: RawData) => {
      const gate = gateInbound(rawToString(data), token);
      if (!gate.ok) {
        log.info("frame-rejected", { cause: gate.cause });
        emit.error(gate.cause, gate.detail);
        if (gate.cause === "unauthorized" || !authenticated)
          ws.close(1008, gate.cause);
        return;
      }

      const message = gate.message;
      if (!authenticated) {
        if (message.type !== "hello") {
          emit.error(
            "invalid-message",
            `handshake required: expected "hello", received "${message.type}"`,
          );
          ws.close(1008, "handshake-required");
          return;
        }
        authenticated = true;
        log.info("handshake-complete", {});
        send({ type: "ready" });
        return;
      }

      if (message.type === "hello") {
        // Idempotent: re-confirm readiness.
        send({ type: "ready" });
        return;
      }
      if (RUNNER_EMITTED_TYPES.has(message.type)) {
        emit.error(
          "invalid-message",
          `"${message.type}" is a runner-emitted event; clients must not send it`,
        );
        return;
      }

      const command = message as CommandMessage;
      log.info("command-received", { command: command.type });
      commandChain = commandChain.then(async () => {
        try {
          await handleCommand(command, {
            emit,
            ready: () => send({ type: "ready" }),
            options,
          });
        } catch (error) {
          // Stages map their own failures; this is the last-resort net.
          log.error("command-crashed", {
            command: command.type,
            error: error instanceof Error ? error.message : String(error),
          });
          emit.error(
            "invalid-message",
            `the "${command.type}" command crashed unexpectedly`,
          );
        }
      });
    });

    ws.on("close", () => {
      log.info("connection-closed", {});
    });
  });

  let closed: Promise<void> | null = null;
  return {
    host,
    port: address.port,
    close: () =>
      (closed ??= (async () => {
        for (const client of wss.clients) client.terminate();
        await new Promise<void>((resolve, reject) => {
          wss.close((error) => (error ? reject(error) : resolve()));
        });
      })()),
  };
}