560 lines
16 KiB
JavaScript
560 lines
16 KiB
JavaScript
import assert from "node:assert/strict";
|
|
import { chmod, mkdir, mkdtemp, realpath, rm, symlink, writeFile } from "node:fs/promises";
|
|
import os from "node:os";
|
|
import path from "node:path";
|
|
import { spawn } from "node:child_process";
|
|
import readline from "node:readline";
|
|
import test from "node:test";
|
|
|
|
const repoRoot = path.resolve(import.meta.dirname, "..");
|
|
const serverPath = path.join(
|
|
repoRoot,
|
|
"plugins",
|
|
"autodl-remote",
|
|
"mcp",
|
|
"server.mjs",
|
|
);
|
|
|
|
class McpClient {
|
|
constructor(child) {
|
|
this.child = child;
|
|
this.nextId = 1;
|
|
this.pending = new Map();
|
|
this.stderr = "";
|
|
|
|
readline.createInterface({ input: child.stdout }).on("line", (line) => {
|
|
const message = JSON.parse(line);
|
|
const pending = this.pending.get(message.id);
|
|
if (!pending) return;
|
|
this.pending.delete(message.id);
|
|
if (message.error) pending.reject(new Error(JSON.stringify(message.error)));
|
|
else pending.resolve(message.result);
|
|
});
|
|
child.stderr.on("data", (chunk) => {
|
|
this.stderr += chunk;
|
|
});
|
|
child.on("exit", (code) => {
|
|
for (const pending of this.pending.values()) {
|
|
pending.reject(
|
|
new Error(`MCP server exited with ${code}: ${this.stderr.trim()}`),
|
|
);
|
|
}
|
|
this.pending.clear();
|
|
});
|
|
}
|
|
|
|
request(method, params = {}) {
|
|
const id = this.nextId++;
|
|
return new Promise((resolve, reject) => {
|
|
this.pending.set(id, { resolve, reject });
|
|
this.child.stdin.write(
|
|
`${JSON.stringify({ jsonrpc: "2.0", id, method, params })}\n`,
|
|
);
|
|
});
|
|
}
|
|
|
|
close() {
|
|
this.child.stdin.end();
|
|
}
|
|
}
|
|
|
|
async function createRecordingCli(tempDir) {
|
|
const fakeCli = path.join(tempDir, "autodl-remote-recording");
|
|
await writeFile(
|
|
fakeCli,
|
|
[
|
|
"#!/usr/bin/env node",
|
|
"console.log(JSON.stringify({ args: process.argv.slice(2), cwd: process.cwd() }));",
|
|
"",
|
|
].join("\n"),
|
|
);
|
|
await chmod(fakeCli, 0o755);
|
|
return fakeCli;
|
|
}
|
|
|
|
function startClient(fakeCli, env = {}) {
|
|
const child = spawn(process.execPath, [serverPath], {
|
|
cwd: repoRoot,
|
|
env: { ...process.env, AUTODL_REMOTE_MCP_CLI: fakeCli, ...env },
|
|
stdio: ["pipe", "pipe", "pipe"],
|
|
});
|
|
return new McpClient(child);
|
|
}
|
|
|
|
async function call(client, name, args) {
|
|
const result = await client.request("tools/call", {
|
|
name,
|
|
arguments: args,
|
|
});
|
|
assert.equal(result.isError, false);
|
|
return JSON.parse(result.structuredContent.stdout);
|
|
}
|
|
|
|
test("MCP exposes account list through the existing CLI", async (t) => {
|
|
const tempDir = await mkdtemp(path.join(os.tmpdir(), "autodl-mcp-"));
|
|
t.after(async () => rm(tempDir, { recursive: true, force: true }));
|
|
|
|
const fakeCli = path.join(tempDir, "autodl-remote");
|
|
await writeFile(
|
|
fakeCli,
|
|
"#!/usr/bin/env bash\nset -euo pipefail\nprintf '* gpu-a\\n gpu-b\\n'\n",
|
|
);
|
|
await chmod(fakeCli, 0o755);
|
|
|
|
const child = spawn(process.execPath, [serverPath], {
|
|
cwd: repoRoot,
|
|
env: { ...process.env, AUTODL_REMOTE_MCP_CLI: fakeCli },
|
|
stdio: ["pipe", "pipe", "pipe"],
|
|
});
|
|
const client = new McpClient(child);
|
|
t.after(() => client.close());
|
|
|
|
const initialized = await client.request("initialize", {
|
|
protocolVersion: "2025-06-18",
|
|
capabilities: {},
|
|
clientInfo: { name: "test", version: "1.0.0" },
|
|
});
|
|
assert.equal(initialized.serverInfo.version, "0.9.0");
|
|
|
|
const listed = await client.request("tools/list");
|
|
assert.ok(listed.tools.some((tool) => tool.name === "autodl_account"));
|
|
|
|
const called = await client.request("tools/call", {
|
|
name: "autodl_account",
|
|
arguments: { action: "list" },
|
|
});
|
|
assert.equal(called.isError, false);
|
|
assert.equal(called.structuredContent.ok, true);
|
|
assert.equal(called.structuredContent.command, "account list");
|
|
assert.match(called.structuredContent.stdout, /gpu-a/);
|
|
});
|
|
|
|
test("MCP maps experiment, fleet, tmux, dashboard, and shutdown workflows", async (t) => {
|
|
const tempDir = await mkdtemp(path.join(os.tmpdir(), "autodl-mcp-workflows-"));
|
|
t.after(async () => rm(tempDir, { recursive: true, force: true }));
|
|
const projectRoot = path.join(tempDir, "project");
|
|
await mkdir(projectRoot);
|
|
const client = startClient(await createRecordingCli(tempDir));
|
|
t.after(() => client.close());
|
|
|
|
const run = await call(client, "autodl_run", {
|
|
action: "submit",
|
|
projectRoot,
|
|
argv: ["python", "train.py"],
|
|
syncUp: ["src", "configs/base.yaml"],
|
|
tmux: true,
|
|
name: "baseline-a",
|
|
model: "qwen",
|
|
tags: ["baseline"],
|
|
});
|
|
assert.deepEqual(run.args, [
|
|
"run",
|
|
"submit",
|
|
"--sync-up",
|
|
"src",
|
|
"--sync-up",
|
|
"configs/base.yaml",
|
|
"--tmux",
|
|
"--name",
|
|
"baseline-a",
|
|
"--model",
|
|
"qwen",
|
|
"--tag",
|
|
"baseline",
|
|
"--",
|
|
"python",
|
|
"train.py",
|
|
]);
|
|
|
|
const jobTail = await call(client, "autodl_job", {
|
|
action: "tail",
|
|
projectRoot,
|
|
name: "baseline-a",
|
|
lines: 40,
|
|
});
|
|
assert.deepEqual(jobTail.args, ["job", "tail", "--lines", "40", "baseline-a"]);
|
|
|
|
const fleet = await call(client, "autodl_fleet", {
|
|
action: "add",
|
|
projectRoot,
|
|
fleet: "lab",
|
|
device: "gpu-1",
|
|
account: "gpu-a",
|
|
remote: "/root/workspace/a",
|
|
tags: "a100,80g",
|
|
});
|
|
assert.deepEqual(fleet.args, [
|
|
"fleet",
|
|
"add",
|
|
"lab",
|
|
"gpu-1",
|
|
"--account",
|
|
"gpu-a",
|
|
"--remote",
|
|
"/root/workspace/a",
|
|
"--tags",
|
|
"a100,80g",
|
|
]);
|
|
|
|
const captured = await call(client, "autodl_tmux", {
|
|
action: "capture",
|
|
projectRoot,
|
|
session: "baseline-a",
|
|
lines: 25,
|
|
});
|
|
assert.deepEqual(captured.args, [
|
|
"tmux",
|
|
"capture",
|
|
"--lines",
|
|
"25",
|
|
"--",
|
|
"baseline-a",
|
|
]);
|
|
|
|
const dashboard = await call(client, "autodl_dashboard", {
|
|
projectRoot,
|
|
fleet: "lab",
|
|
output: "reports/fleet.html",
|
|
lines: 60,
|
|
});
|
|
assert.deepEqual(dashboard.args, [
|
|
"dashboard",
|
|
"--fleet",
|
|
"lab",
|
|
"--output",
|
|
"reports/fleet.html",
|
|
"--lines",
|
|
"60",
|
|
]);
|
|
|
|
const shutdown = await call(client, "autodl_shutdown", {
|
|
projectRoot,
|
|
confirm: "shutdown",
|
|
waitSeconds: 3,
|
|
});
|
|
assert.deepEqual(shutdown.args, ["shutdown", "--wait", "3"]);
|
|
});
|
|
|
|
test("MCP requires explicit confirmation for destructive operations", async (t) => {
|
|
const tempDir = await mkdtemp(path.join(os.tmpdir(), "autodl-mcp-safety-"));
|
|
t.after(async () => rm(tempDir, { recursive: true, force: true }));
|
|
const projectRoot = path.join(tempDir, "project");
|
|
await mkdir(projectRoot);
|
|
const client = startClient(await createRecordingCli(tempDir));
|
|
t.after(() => client.close());
|
|
|
|
await assert.rejects(
|
|
client.request("tools/call", {
|
|
name: "autodl_job",
|
|
arguments: { action: "stop", projectRoot, name: "train-a", confirm: "train-b" },
|
|
}),
|
|
/exactly match name/,
|
|
);
|
|
await assert.rejects(
|
|
client.request("tools/call", {
|
|
name: "autodl_exec",
|
|
arguments: {
|
|
projectRoot,
|
|
argv: ["shutdown"],
|
|
allowDangerous: true,
|
|
},
|
|
}),
|
|
/confirmDangerous=allow-dangerous/,
|
|
);
|
|
await assert.rejects(
|
|
client.request("tools/call", {
|
|
name: "autodl_transfer",
|
|
arguments: {
|
|
action: "upload",
|
|
projectRoot,
|
|
localPath: "src",
|
|
delete: true,
|
|
},
|
|
}),
|
|
/confirmDelete=delete-remote-extraneous/,
|
|
);
|
|
await assert.rejects(
|
|
client.request("tools/call", {
|
|
name: "autodl_file",
|
|
arguments: { action: "read", projectRoot },
|
|
}),
|
|
/path is required/,
|
|
);
|
|
await assert.rejects(
|
|
client.request("tools/call", {
|
|
name: "autodl_tmux",
|
|
arguments: { action: "install", projectRoot },
|
|
}),
|
|
/confirm=install-tmux/,
|
|
);
|
|
await assert.rejects(
|
|
client.request("tools/call", {
|
|
name: "autodl_shutdown",
|
|
arguments: { projectRoot, confirm: "yes" },
|
|
}),
|
|
/confirm=shutdown/,
|
|
);
|
|
});
|
|
|
|
test("MCP keeps declared file operands inside the selected project", async (t) => {
|
|
const tempDir = await mkdtemp(path.join(os.tmpdir(), "autodl-mcp-boundary-"));
|
|
t.after(async () => rm(tempDir, { recursive: true, force: true }));
|
|
const projectRoot = path.join(tempDir, "project");
|
|
await mkdir(projectRoot);
|
|
const outside = path.join(tempDir, "outside.txt");
|
|
await writeFile(outside, "secret");
|
|
await symlink(outside, path.join(projectRoot, "outside-link"));
|
|
const client = startClient(await createRecordingCli(tempDir));
|
|
t.after(() => client.close());
|
|
|
|
const rejectedCalls = [
|
|
["autodl_file", { action: "read", projectRoot, path: "/etc/passwd" }],
|
|
["autodl_file", { action: "tail", projectRoot, path: "../outside.txt" }],
|
|
["autodl_transfer", { action: "upload", projectRoot, localPath: "/etc/passwd" }],
|
|
["autodl_transfer", { action: "upload", projectRoot, localPath: "outside-link" }],
|
|
["autodl_transfer", { action: "download", projectRoot, remotePath: "runs/a", localPath: "../outside.txt" }],
|
|
["autodl_transfer", { action: "download", projectRoot, remotePath: "/etc/passwd", localPath: "copy" }],
|
|
["autodl_exec", { projectRoot, cwd: "../outside", argv: ["pwd"] }],
|
|
["autodl_run", { action: "submit", projectRoot, syncUp: ["../outside.txt"], argv: ["true"] }],
|
|
["autodl_dashboard", { projectRoot, output: "../outside.html" }],
|
|
];
|
|
for (const [name, arguments_] of rejectedCalls) {
|
|
await assert.rejects(
|
|
client.request("tools/call", { name, arguments: arguments_ }),
|
|
/project-relative|escape projectRoot/,
|
|
);
|
|
}
|
|
});
|
|
|
|
test("MCP rejects ambiguous names and negotiates a supported protocol version", async (t) => {
|
|
const tempDir = await mkdtemp(path.join(os.tmpdir(), "autodl-mcp-canonical-"));
|
|
t.after(async () => rm(tempDir, { recursive: true, force: true }));
|
|
const projectRoot = path.join(tempDir, "project");
|
|
await mkdir(projectRoot);
|
|
const client = startClient(await createRecordingCli(tempDir));
|
|
t.after(() => client.close());
|
|
|
|
await assert.rejects(
|
|
client.request("tools/call", {
|
|
name: "autodl_job",
|
|
arguments: { action: "stop", projectRoot, name: "train prod", confirm: "train prod" },
|
|
}),
|
|
/canonical name/,
|
|
);
|
|
const initialized = await client.request("initialize", {
|
|
protocolVersion: "2099-01-01",
|
|
capabilities: {},
|
|
clientInfo: { name: "future", version: "1" },
|
|
});
|
|
assert.equal(initialized.protocolVersion, "2025-06-18");
|
|
});
|
|
|
|
test("MCP returns bounded structured output for CLI failures", async (t) => {
|
|
const tempDir = await mkdtemp(path.join(os.tmpdir(), "autodl-mcp-output-"));
|
|
t.after(async () => rm(tempDir, { recursive: true, force: true }));
|
|
const projectRoot = path.join(tempDir, "project");
|
|
await mkdir(projectRoot);
|
|
const fakeCli = path.join(tempDir, "autodl-remote-failing");
|
|
await writeFile(
|
|
fakeCli,
|
|
[
|
|
"#!/usr/bin/env node",
|
|
"process.stdout.write('A'.repeat(100));",
|
|
"process.stderr.write('B'.repeat(100));",
|
|
"process.exitCode = 7;",
|
|
"",
|
|
].join("\n"),
|
|
);
|
|
await chmod(fakeCli, 0o755);
|
|
const client = startClient(fakeCli, { AUTODL_REMOTE_MCP_MAX_OUTPUT_BYTES: "16" });
|
|
t.after(() => client.close());
|
|
|
|
const result = await client.request("tools/call", {
|
|
name: "autodl_doctor",
|
|
arguments: { projectRoot },
|
|
});
|
|
assert.equal(result.isError, true);
|
|
assert.equal(result.structuredContent.ok, false);
|
|
assert.equal(result.structuredContent.exitCode, 7);
|
|
assert.equal(Buffer.byteLength(result.structuredContent.stdout), 16);
|
|
assert.equal(Buffer.byteLength(result.structuredContent.stderr), 16);
|
|
assert.equal(result.structuredContent.stdoutTruncated, true);
|
|
assert.equal(result.structuredContent.stderrTruncated, true);
|
|
});
|
|
|
|
test("MCP drains an in-flight CLI call after stdin reaches EOF", async (t) => {
|
|
const tempDir = await mkdtemp(path.join(os.tmpdir(), "autodl-mcp-eof-"));
|
|
t.after(async () => rm(tempDir, { recursive: true, force: true }));
|
|
const fakeCli = path.join(tempDir, "autodl-remote-slow");
|
|
await writeFile(
|
|
fakeCli,
|
|
[
|
|
"#!/usr/bin/env node",
|
|
"setTimeout(() => console.log('drained'), 100);",
|
|
"",
|
|
].join("\n"),
|
|
);
|
|
await chmod(fakeCli, 0o755);
|
|
|
|
const child = spawn(process.execPath, [serverPath], {
|
|
cwd: repoRoot,
|
|
env: { ...process.env, AUTODL_REMOTE_MCP_CLI: fakeCli },
|
|
stdio: ["pipe", "pipe", "pipe"],
|
|
});
|
|
const lines = [];
|
|
readline.createInterface({ input: child.stdout }).on("line", (line) => {
|
|
lines.push(JSON.parse(line));
|
|
});
|
|
let stderr = "";
|
|
child.stderr.on("data", (chunk) => {
|
|
stderr += chunk;
|
|
});
|
|
child.stdin.end(
|
|
`${JSON.stringify({
|
|
jsonrpc: "2.0",
|
|
id: 1,
|
|
method: "tools/call",
|
|
params: { name: "autodl_account", arguments: { action: "list" } },
|
|
})}\n`,
|
|
);
|
|
const exitCode = await new Promise((resolve) => child.on("exit", resolve));
|
|
assert.equal(exitCode, 0, stderr);
|
|
assert.equal(lines.length, 1);
|
|
assert.equal(lines[0].result.structuredContent.ok, true);
|
|
assert.match(lines[0].result.structuredContent.stdout, /drained/);
|
|
});
|
|
|
|
test("MCP exits cleanly when the client closes its response stream", async (t) => {
|
|
const tempDir = await mkdtemp(path.join(os.tmpdir(), "autodl-mcp-epipe-"));
|
|
t.after(async () => rm(tempDir, { recursive: true, force: true }));
|
|
const child = spawn(process.execPath, [serverPath], {
|
|
cwd: repoRoot,
|
|
env: {
|
|
...process.env,
|
|
AUTODL_REMOTE_MCP_CLI: await createRecordingCli(tempDir),
|
|
},
|
|
stdio: ["pipe", "pipe", "pipe"],
|
|
});
|
|
let stderr = "";
|
|
child.stderr.on("data", (chunk) => {
|
|
stderr += chunk;
|
|
});
|
|
child.stdout.destroy();
|
|
child.stdin.end(
|
|
`${JSON.stringify({
|
|
jsonrpc: "2.0",
|
|
id: 1,
|
|
method: "tools/call",
|
|
params: { name: "autodl_account", arguments: { action: "list" } },
|
|
})}\n`,
|
|
);
|
|
const exitCode = await new Promise((resolve) => child.on("exit", resolve));
|
|
assert.equal(exitCode, 0, stderr);
|
|
assert.doesNotMatch(stderr, /EPIPE/);
|
|
});
|
|
|
|
test("MCP standardizes common project, file, transfer, and execution commands", async (t) => {
|
|
const tempDir = await mkdtemp(path.join(os.tmpdir(), "autodl-mcp-map-"));
|
|
t.after(async () => rm(tempDir, { recursive: true, force: true }));
|
|
const projectRoot = path.join(tempDir, "project");
|
|
await mkdir(projectRoot);
|
|
const client = startClient(await createRecordingCli(tempDir));
|
|
t.after(() => client.close());
|
|
|
|
const listed = await client.request("tools/list");
|
|
assert.deepEqual(
|
|
listed.tools.map((tool) => tool.name),
|
|
[
|
|
"autodl_account",
|
|
"autodl_bind",
|
|
"autodl_doctor",
|
|
"autodl_model_dir",
|
|
"autodl_file",
|
|
"autodl_transfer",
|
|
"autodl_exec",
|
|
"autodl_job",
|
|
"autodl_run",
|
|
"autodl_fleet",
|
|
"autodl_tmux",
|
|
"autodl_dashboard",
|
|
"autodl_shutdown",
|
|
],
|
|
);
|
|
|
|
const bound = await call(client, "autodl_bind", {
|
|
projectRoot,
|
|
account: "gpu-a",
|
|
remote: "/root/workspace/demo",
|
|
force: true,
|
|
});
|
|
assert.deepEqual(bound.args, [
|
|
"bind",
|
|
"--account",
|
|
"gpu-a",
|
|
"--remote",
|
|
"/root/workspace/demo",
|
|
"--force",
|
|
]);
|
|
assert.equal(bound.cwd, await realpath(projectRoot));
|
|
|
|
const read = await call(client, "autodl_file", {
|
|
action: "read",
|
|
projectRoot,
|
|
path: "logs/train.log",
|
|
lines: 80,
|
|
account: "gpu-a",
|
|
});
|
|
assert.deepEqual(read.args, [
|
|
"cat",
|
|
"--account",
|
|
"gpu-a",
|
|
"--lines",
|
|
"80",
|
|
"--",
|
|
"logs/train.log",
|
|
]);
|
|
|
|
const uploaded = await call(client, "autodl_transfer", {
|
|
action: "upload",
|
|
projectRoot,
|
|
localPath: "src",
|
|
remotePath: "src",
|
|
dryRun: true,
|
|
delete: true,
|
|
confirmDelete: "delete-remote-extraneous",
|
|
});
|
|
assert.deepEqual(uploaded.args, [
|
|
"put",
|
|
"--dry-run",
|
|
"--delete",
|
|
"--",
|
|
"src",
|
|
"src",
|
|
]);
|
|
|
|
const executed = await call(client, "autodl_exec", {
|
|
projectRoot,
|
|
argv: ["bash", "-lc", "python train.py | tee train.log"],
|
|
mode: "tmux",
|
|
name: "train-a",
|
|
tags: ["baseline", "a100"],
|
|
outputs: ["runs/train-a"],
|
|
});
|
|
assert.deepEqual(executed.args, [
|
|
"exec",
|
|
"--tmux",
|
|
"--name",
|
|
"train-a",
|
|
"--tag",
|
|
"baseline",
|
|
"--tag",
|
|
"a100",
|
|
"--output",
|
|
"runs/train-a",
|
|
"--",
|
|
"bash",
|
|
"-lc",
|
|
"python train.py | tee train.log",
|
|
]);
|
|
});
|