agentic-dev/apps/runner/tests/protocol-client.ts

import { once } from "node:events";
import WebSocket from "ws";
import {
  envelopeSchema,
  PROTOCOL_VERSION,
  type Envelope,
  type RunnerMessage,
} from "@repo/core-runner-protocol";

/**
 * Test-side protocol client. Every inbound frame is parsed against the
 * envelope schema — so every assertion made through this client also
 * proves the runner's outbound messages are protocol-conformant.
 */
export interface ProtocolClient {
  /** Envelope-wrap and send a message (token/version overridable per call). */
  send: (
    message: RunnerMessage,
    overrides?: Partial<Omit<Envelope, "message">>,
  ) => void;
  /** Send a raw frame verbatim (malformed-frame tests). */
  sendRaw: (data: string) => void;
  /** All inbound messages so far, in arrival order. */
  received: () => RunnerMessage[];
  /** All inbound frames verbatim (for leak assertions on the raw wire). */
  rawReceived: () => string[];
  /** Resolve the first (possibly already-received) message matching. */
  waitFor: (
    predicate: (message: RunnerMessage) => boolean,
    description: string,
    timeoutMs?: number,
  ) => Promise<RunnerMessage>;
  /** Resolves with the close code once the server closes the socket. */
  closed: Promise<number>;
  close: () => void;
}

export async function connectProtocolClient(
  port: number,
  token: string,
): Promise<ProtocolClient> {
  const ws = new WebSocket(`ws://127.0.0.1:${port}`);
  await once(ws, "open");

  const inbound: RunnerMessage[] = [];
  const rawFrames: string[] = [];
  const waiters = new Set<() => void>();

  ws.on("message", (data) => {
    const raw = String(data);
    rawFrames.push(raw);
    const envelope = envelopeSchema.parse(JSON.parse(raw));
    inbound.push(envelope.message);
    for (const notify of [...waiters]) notify();
  });

  const closed = new Promise<number>((resolve) => {
    ws.on("close", (code) => resolve(code));
  });

  return {
    send: (message, overrides) => {
      ws.send(
        JSON.stringify({
          protocolVersion: PROTOCOL_VERSION,
          token,
          ...overrides,
          message,
        }),
      );
    },
    sendRaw: (data) => ws.send(data),
    received: () => [...inbound],
    rawReceived: () => [...rawFrames],
    waitFor: (predicate, description, timeoutMs = 30_000) => {
      return new Promise<RunnerMessage>((resolve, reject) => {
        let scanned = 0;
        const scan = (): void => {
          while (scanned < inbound.length) {
            const message = inbound[scanned];
            scanned += 1;
            if (message !== undefined && predicate(message)) {
              cleanup();
              resolve(message);
              return;
            }
          }
        };
        const timer = setTimeout(() => {
          cleanup();
          reject(
            new Error(
              `timed out waiting for ${description}; received: ${JSON.stringify(inbound.map((m) => m.type))}`,
            ),
          );
        }, timeoutMs);
        const cleanup = (): void => {
          clearTimeout(timer);
          waiters.delete(scan);
        };
        waiters.add(scan);
        scan();
      });
    },
    closed,
    close: () => ws.close(),
  };
}

/** Complete the hello/ready handshake and return the ready message. */
export async function handshake(
  client: ProtocolClient,
): Promise<RunnerMessage> {
  client.send({ type: "hello" });
  return client.waitFor((m) => m.type === "ready", "ready handshake reply");
}

/**
 * Send a command and resolve on the runner's NEXT `ready` — the
 * protocol's success marker for a completed command. Counting readies
 * (instead of matching any `ready`) keeps this correct on sessions that
 * already completed earlier commands.
 */
export async function sendAndAwaitReady(
  client: ProtocolClient,
  message: RunnerMessage,
  description: string,
  timeoutMs?: number,
): Promise<void> {
  const readiesBefore = client
    .received()
    .filter((m) => m.type === "ready").length;
  client.send(message);
  await client.waitFor(
    (m) =>
      m.type === "ready" &&
      client.received().filter((r) => r.type === "ready").length >
        readiesBefore,
    description,
    timeoutMs,
  );
}

/** Wait for the next `error` event and assert its named cause. */
export async function expectNamedError(
  client: ProtocolClient,
  cause: string,
): Promise<Extract<RunnerMessage, { type: "error" }>> {
  const error = await client.waitFor(
    (m) => m.type === "error",
    `${cause} error`,
  );
  if (error.type !== "error" || error.cause !== cause) {
    throw new Error(
      `expected a "${cause}" error, received: ${JSON.stringify(error)}`,
    );
  }
  return error;
}