agentic-dev/apps/runner/tests/protocol-server.test.ts

/**
 * Protocol-server suite over a real WS connection (in-process server so
 * v8 coverage sees `server.ts`; the spawn suite in
 * `runner-process.integration.test.ts` proves the same handshake against
 * a real child process). Assertions go through the protocol only — every
 * inbound frame the client records is envelope-parsed, so runner
 * outbound conformance is asserted implicitly on every test.
 */
import { mkdtemp, rm } from "node:fs/promises";
import os from "node:os";
import path from "node:path";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { PROTOCOL_VERSION } from "@repo/core-runner-protocol";
import { startRunnerServer, type RunnerServer } from "@/server";
import { createLogger } from "@/log";
import {
  connectProtocolClient,
  expectNamedError,
  handshake,
  type ProtocolClient,
} from "./protocol-client";

const TOKEN = "test-workspace-token";

let server: RunnerServer;
let workspaceDir: string;
let logLines: string[];
let clients: ProtocolClient[];

async function connect(token = TOKEN): Promise<ProtocolClient> {
  const client = await connectProtocolClient(server.port, token);
  clients.push(client);
  return client;
}

beforeEach(async () => {
  workspaceDir = await mkdtemp(path.join(os.tmpdir(), "veect-runner-test-"));
  logLines = [];
  clients = [];
  server = await startRunnerServer({
    host: "127.0.0.1",
    port: 0,
    token: TOKEN,
    workspaceDir,
    heartbeatMs: 50,
    log: createLogger((line) => logLines.push(line)),
  });
});

afterEach(async () => {
  for (const client of clients) client.close();
  await server.close();
  await rm(workspaceDir, { recursive: true, force: true });
});

describe("hello/ready handshake", () => {
  it("answers a valid hello with ready", async () => {
    const client = await connect();
    const ready = await handshake(client);
    expect(ready).toEqual({ type: "ready" });
  });

  it.each([
    [
      "a bad token",
      "unauthorized",
      (client: ProtocolClient) => client.send({ type: "hello" }),
      "wrong-token",
    ],
    [
      "a missing token (schema)",
      "invalid-message",
      (client: ProtocolClient) =>
        client.sendRaw(
          JSON.stringify({
            protocolVersion: PROTOCOL_VERSION,
            message: { type: "hello" },
          }),
        ),
      TOKEN,
    ],
    [
      "an unsupported protocol version",
      "unsupported-protocol-version",
      (client: ProtocolClient) =>
        client.send({ type: "hello" }, { protocolVersion: "42" as never }),
      TOKEN,
    ],
    [
      "a non-JSON frame",
      "invalid-message",
      (client: ProtocolClient) => client.sendRaw("definitely not json"),
      TOKEN,
    ],
  ] as const)(
    "rejects %s with a named %s error and closes 1008",
    async (_label, cause, act, token) => {
      const client = await connect(token);
      act(client);
      await expectNamedError(client, cause);
      await expect(client.closed).resolves.toBe(1008);
    },
  );

  it("requires hello first: a command before the handshake is rejected", async () => {
    const client = await connect();
    client.send({ type: "install" });
    const error = await expectNamedError(client, "invalid-message");
    expect(error.message).toContain("handshake required");
    await expect(client.closed).resolves.toBe(1008);
  });

  it("redacts the workspace token on replies to unauthenticated peers", async () => {
    const client = await connect("wrong-token");
    client.send({ type: "hello" });
    await expectNamedError(client, "unauthorized");
    // The rejection envelope must not carry the real workspace token.
    expect(client.rawReceived().join("")).not.toContain(TOKEN);
    expect(client.rawReceived().join("")).toContain('"redacted"');
    await expect(client.closed).resolves.toBe(1008);
  });

  it("answers a repeated hello idempotently with ready", async () => {
    const client = await connect();
    await handshake(client);
    client.send({ type: "hello" });
    const second = await client.waitFor(
      (m) =>
        m.type === "ready" &&
        client.received().filter((r) => r.type === "ready").length >= 2,
      "second ready",
    );
    expect(second).toEqual({ type: "ready" });
  });
});

describe("post-handshake message gate", () => {
  it("keeps the session open after a malformed frame post-handshake", async () => {
    const client = await connect();
    await handshake(client);
    client.sendRaw("{broken");
    await expectNamedError(client, "invalid-message");
    // Session survives: another hello still gets a ready.
    client.send({ type: "hello" });
    const ready = await client.waitFor(
      (m) =>
        m.type === "ready" &&
        client.received().filter((r) => r.type === "ready").length >= 2,
      "ready after malformed frame",
    );
    expect(ready).toEqual({ type: "ready" });
  });

  it("closes the session on a bad token even post-handshake", async () => {
    const client = await connect();
    await handshake(client);
    client.send({ type: "install" }, { token: "tampered" });
    await expectNamedError(client, "unauthorized");
    await expect(client.closed).resolves.toBe(1008);
  });

  it("rejects runner-emitted event types sent by the client", async () => {
    const client = await connect();
    await handshake(client);
    client.send({ type: "status", stage: "cloning", elapsedMs: 1 });
    const error = await expectNamedError(client, "invalid-message");
    expect(error.message).toContain("runner-emitted");
  });

  it("rejects a strict-schema violation without echoing payload values", async () => {
    const client = await connect();
    await handshake(client);
    client.sendRaw(
      JSON.stringify({
        protocolVersion: PROTOCOL_VERSION,
        token: TOKEN,
        message: {
          type: "clone",
          gitUrl: "git://x/y.git",
          pat: "sekrit-pat",
          extra: true,
        },
      }),
    );
    const error = await expectNamedError(client, "invalid-message");
    expect(JSON.stringify(error)).not.toContain("sekrit-pat");
    expect(logLines.join("\n")).not.toContain("sekrit-pat");
  });
});

describe("commands not yet wired (stories 05+)", () => {
  it.each([
    ["scan", "scan-failed", "scanning"],
    ["adapter-start", "adapter-start-failed", "starting-preview"],
  ] as const)(
    "answers %s with the named %s error",
    async (type, cause, stage) => {
      const client = await connect();
      await handshake(client);
      client.send({ type } as never);
      const error = await expectNamedError(client, cause);
      expect(error).toMatchObject({ type: "error", cause, stage });
    },
  );

  it("answers render-frame with the named render-failed error", async () => {
    const client = await connect();
    await handshake(client);
    client.send({ type: "render-frame", componentId: "button" });
    await expectNamedError(client, "render-failed");
  });
});