cc-libs/tools/mcp-bridge/test-integration/harness.ts

import { spawn } from "node:child_process";
import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import type { Server } from "node:http";
import type { WebSocketServer } from "ws";
import { LinkRegistry, startLinkServer } from "../src/link-server.js";
import { startMcpServer } from "../src/mcp-server.js";

const HERE = dirname(fileURLToPath(import.meta.url));
const LUA_DIR = join(HERE, "lua");
const REPO_ROOT = join(HERE, "../../..");

export type Bridge = {
  registry: LinkRegistry;
  mcpUrl: string;
  linkUrl: string;
  close: () => Promise<void>;
};

export async function startBridge(probeTimeoutMs = 500): Promise<Bridge> {
  const registry = new LinkRegistry();
  const mcpServer = startMcpServer({ host: "127.0.0.1", port: 0, probeTimeoutMs, registry });
  const linkServer = startLinkServer({ host: "127.0.0.1", port: 0, registry });
  await Promise.all([waitForListening(mcpServer), waitForListening(linkServer)]);
  return {
    registry,
    mcpUrl: `http://127.0.0.1:${serverPort(mcpServer)}`,
    linkUrl: `ws://127.0.0.1:${serverPort(linkServer)}`,
    close: () => closeBridge(mcpServer, linkServer),
  };
}

export async function waitForComputers(registry: LinkRegistry, count: number, timeoutMs: number): Promise<void> {
  const deadline = Date.now() + timeoutMs;
  while (Date.now() < deadline) {
    if (registry.count() >= count) {
      return;
    }
    await sleep(50);
  }
  throw new Error(`waitForComputers timed out (expected ${count}, got ${registry.count()})`);
}

export async function callProbeComputers(mcpUrl: string): Promise<string> {
  const body = {
    jsonrpc: "2.0",
    id: 1,
    method: "tools/call",
    params: { name: "probe-computers", arguments: {} },
  };
  const response = await fetch(mcpUrl, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify(body),
  });
  const payload = (await response.json()) as {
    result?: { content?: { type: string; text: string }[] };
  };
  const text = payload.result?.content?.[0]?.text;
  if (typeof text !== "string") {
    throw new Error(`Unexpected MCP response: ${JSON.stringify(payload)}`);
  }
  return text;
}

export async function callExecLua(mcpUrl: string, computerId: number, code: string, timeoutMs?: number): Promise<string> {
  const body = {
    jsonrpc: "2.0",
    id: 1,
    method: "tools/call",
    params: { name: "exec-lua", arguments: { computerId, code, timeoutMs } },
  };
  const response = await fetch(mcpUrl, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify(body),
  });
  const payload = (await response.json()) as {
    result?: { content?: { type: string; text: string }[] };
  };
  const text = payload.result?.content?.[0]?.text;
  if (typeof text !== "string") {
    throw new Error(`Unexpected MCP response: ${JSON.stringify(payload)}`);
  }
  return text;
}

export type CraftosResult = {
  status: number | null;
  signal: NodeJS.Signals | null;
  output: string;
};

export type CraftosHandle = {
  done: Promise<CraftosResult>;
  abort: () => void;
};

export function startCraftos(
  luaName: string,
  opts: { computerId?: number; computerLabel?: string; mountRepo?: boolean; shellArgs?: string[]; timeoutMs?: number } = {},
): CraftosHandle {
  const timeoutMs = opts.timeoutMs ?? 15_000;
  const controller = new AbortController();

  const done = (async (): Promise<CraftosResult> => {
    const dataDir = await mkdtemp(join(tmpdir(), "mcp-bridge-it-"));
    const watchdog = setTimeout(() => controller.abort(), timeoutMs);
    try {
      const args: string[] = ["--directory", dataDir, "--headless", "--mount-ro", `/staging=${LUA_DIR}`];
      if (opts.computerId !== undefined) {
        args.push("--id", String(opts.computerId));
      }
      if (opts.mountRepo) {
        args.push(
          "--mount-ro", `/trapos=${REPO_ROOT}`,
          "--mount-ro", `/apis=${join(REPO_ROOT, "apis")}`,
          "--mount-ro", `/programs=${join(REPO_ROOT, "programs")}`,
        );
      }
      if (process.platform === "darwin") {
        args.push("--rom", "/Applications/CraftOS-PC.app/Contents/Resources");
      }
      args.push("--exec", buildExecCode(luaName, opts.shellArgs ?? [], opts.computerLabel));

      const chunks: Buffer[] = [];
      const child = spawn("craftos", args, { signal: controller.signal });
      child.stdout.on("data", (d: Buffer) => chunks.push(d));
      child.stderr.on("data", (d: Buffer) => chunks.push(d));

      const result = await new Promise<{ status: number | null; signal: NodeJS.Signals | null }>((resolve) => {
        child.once("close", (code, signal) => resolve({ status: code, signal }));
        child.once("error", () => resolve({ status: null, signal: null }));
      });

      return { status: result.status, signal: result.signal, output: Buffer.concat(chunks).toString("utf8") };
    } finally {
      clearTimeout(watchdog);
      await rm(dataDir, { recursive: true, force: true });
    }
  })();

  return { done, abort: () => controller.abort() };
}

export function formatFailure(message: string, craftosOutput: string): string {
  const lines = ["\x1b[31mFAIL\x1b[0m " + message, "--- craftos output ---", craftosOutput.trimEnd(), "----------------------"];
  return lines.join("\n");
}

function buildExecCode(luaName: string, shellArgs: string[], computerLabel?: string): string {
  const programPath = luaName.startsWith("/") ? luaName : `/staging/${luaName}`;
  const parts = [luaQuote(programPath), ...shellArgs.map(luaQuote)];
  const setup = computerLabel === undefined ? "" : `os.setComputerLabel(${luaQuote(computerLabel)}); `;
  return `${setup}shell.run(${parts.join(", ")})`;
}

function luaQuote(value: string): string {
  return `'${value.replaceAll("\\", "\\\\").replaceAll("'", "\\'")}'`;
}

function waitForListening(server: Server | WebSocketServer): Promise<void> {
  return new Promise((resolve, reject) => {
    const addr = "address" in server ? server.address() : null;
    if (addr) {
      resolve();
      return;
    }
    server.once("listening", () => resolve());
    server.once("error", reject);
  });
}

function serverPort(server: Server | WebSocketServer): number {
  const address = server.address();
  if (!address || typeof address === "string") {
    throw new Error("server does not have a TCP address");
  }
  return address.port;
}

async function closeBridge(mcpServer: Server, linkServer: WebSocketServer): Promise<void> {
  for (const client of linkServer.clients) {
    client.terminate();
  }
  await new Promise<void>((resolve) => linkServer.close(() => resolve()));
  await new Promise<void>((resolve) => mcpServer.close(() => resolve()));
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}