feat: close AgentRun commander task plane gaps

2026-06-09 00:06:53 +08:00
parent 96c8283574
commit 2f5cf8b3d4
15 changed files with 411 additions and 23 deletions
@@ -0,0 +1,270 @@
+#!/usr/bin/env bun
+
+const defaultFrontendPort = 18081;
+const sshInputChunkBytes = 32 * 1024;
+
+function jsonHelp() {
+  return {
+    ok: true,
+    tool: "tran",
+    purpose: "AgentRun runner UniDesk SSH passthrough over frontend /ws/ssh",
+    requiredEnv: ["UNIDESK_SSH_CLIENT_TOKEN", "UNIDESK_MAIN_SERVER_IP or UNIDESK_FRONTEND_URL"],
+    supported: [
+      "tran <provider> <command...>",
+      "tran <provider> argv <command...>",
+      "tran <provider> script -- '<shell script>'",
+      "tran <provider>:/absolute/workspace script -- '<shell script>'",
+      "tran <provider>:k3s kubectl <args...>",
+      "tran <provider>:k3s script -- '<shell script>'",
+      "tran <provider>:k3s:<namespace>:<workload>[:container] argv <command...>",
+      "tran <provider>:k3s:<namespace>:<workload>[:container] script -- '<shell script>'",
+    ],
+    unsupported: ["apply-patch", "upload", "download", "Windows win/ps/cmd routes"],
+    valuesPrinted: false,
+  };
+}
+
+function writeJson(value, stream = process.stdout) {
+  stream.write(`${JSON.stringify(value)}\n`);
+}
+
+function fail(failureKind, message, details = {}, exitCode = 2) {
+  writeJson({ ok: false, failureKind, message, ...details, valuesPrinted: false }, process.stderr);
+  process.exit(exitCode);
+}
+
+function shellQuote(value) {
+  return `'${String(value).replace(/'/g, `'"'"'`)}'`;
+}
+
+function shellArgv(args) {
+  if (args.length === 0) return "";
+  return args.map(shellQuote).join(" ");
+}
+
+function baseUrlFromEnv(env) {
+  const explicit = (env.UNIDESK_FRONTEND_URL || env.UNIDESK_MAIN_SERVER_URL || "").trim();
+  if (explicit) return explicit.replace(/\/+$/g, "");
+  const host = (env.UNIDESK_MAIN_SERVER_IP || env.UNIDESK_MAIN_SERVER_HOST || env.CODE_QUEUE_DEV_CONTAINER_MASTER_HOST || "").trim();
+  if (!host) return null;
+  if (host.startsWith("http://") || host.startsWith("https://")) return host.replace(/\/+$/g, "");
+  if (/:[0-9]+$/u.test(host)) return `http://${host}`;
+  const port = Number(env.UNIDESK_FRONTEND_PORT || env.UNIDESK_MAIN_SERVER_PORT || defaultFrontendPort);
+  return `http://${host}:${Number.isInteger(port) && port > 0 ? port : defaultFrontendPort}`;
+}
+
+function websocketUrl(baseUrl) {
+  const url = new URL("/ws/ssh", baseUrl);
+  url.protocol = url.protocol === "https:" ? "wss:" : "ws:";
+  return url.toString();
+}
+
+function parseRoute(raw) {
+  const parts = String(raw).split(":");
+  const providerId = parts.shift() || "";
+  if (!providerId) fail("schema-invalid", "tran route requires a provider id", { route: raw });
+  if (parts.length === 0) return { providerId, plane: "host", workspace: null, namespace: null, resource: null, container: null, raw };
+  if (parts[0] === "win") fail("unsupported-operation", "AgentRun tran does not support Windows routes yet", { route: raw });
+  if (parts[0] === "k3s") {
+    return {
+      providerId,
+      plane: "k3s",
+      workspace: null,
+      namespace: parts[1] || null,
+      resource: parts[2] || null,
+      container: parts[3] || null,
+      raw,
+    };
+  }
+  const workspace = parts.join(":");
+  if (!workspace.startsWith("/")) fail("schema-invalid", "host workspace routes must be absolute paths", { route: raw });
+  return { providerId, plane: "host", workspace, namespace: null, resource: null, container: null, raw };
+}
+
+async function readStdinText() {
+  const chunks = [];
+  for await (const chunk of Bun.stdin.stream()) chunks.push(Buffer.from(chunk));
+  return Buffer.concat(chunks).toString("utf8");
+}
+
+async function scriptCommand(args) {
+  if (args[0] === "--") {
+    const rest = args.slice(1);
+    if (rest.length === 0) return await readStdinText();
+    if (rest.length === 1) return rest[0];
+    return shellArgv(rest);
+  }
+  if (args.length === 0) return await readStdinText();
+  return shellArgv(args);
+}
+
+async function hostCommand(route, args) {
+  if (args.length === 0) return { command: null, cwd: route.workspace, tty: true };
+  const op = args[0];
+  if (op === "apply-patch" || op === "upload" || op === "download") {
+    fail("unsupported-operation", `AgentRun tran does not support ${op}; use host/source controlled tools outside the runner for that operation`, { operation: op });
+  }
+  if (op === "script" || op === "shell") return { command: await scriptCommand(args.slice(1)), cwd: route.workspace, tty: false };
+  if (op === "argv") return { command: shellArgv(args.slice(1)), cwd: route.workspace, tty: false };
+  return { command: shellArgv(args), cwd: route.workspace, tty: false };
+}
+
+function k3sExecPrefix(route) {
+  const base = ["env", "KUBECONFIG=/etc/rancher/k3s/k3s.yaml", "kubectl", "exec"];
+  if (route.namespace) base.push("-n", route.namespace);
+  if (!route.resource) fail("schema-invalid", "k3s workload routes require namespace and resource", { route: route.raw });
+  base.push(route.resource);
+  if (route.container) base.push("-c", route.container);
+  base.push("--");
+  return base;
+}
+
+async function k3sCommand(route, args) {
+  const op = args[0] || "kubectl";
+  if (op === "apply-patch" || op === "upload" || op === "download") {
+    fail("unsupported-operation", `AgentRun tran does not support ${op}; use host/source controlled tools outside the runner for that operation`, { operation: op });
+  }
+  if (!route.resource) {
+    if (op === "kubectl") return { command: shellArgv(["env", "KUBECONFIG=/etc/rancher/k3s/k3s.yaml", "kubectl", ...args.slice(1)]), cwd: null, tty: false };
+    if (op === "script" || op === "shell") {
+      const script = await scriptCommand(args.slice(1));
+      return { command: `export KUBECONFIG=/etc/rancher/k3s/k3s.yaml; ${script}`, cwd: null, tty: false };
+    }
+    if (op === "argv") return { command: shellArgv(["env", "KUBECONFIG=/etc/rancher/k3s/k3s.yaml", ...args.slice(1)]), cwd: null, tty: false };
+    return { command: shellArgv(["env", "KUBECONFIG=/etc/rancher/k3s/k3s.yaml", ...args]), cwd: null, tty: false };
+  }
+  if (op === "script" || op === "shell") {
+    const script = await scriptCommand(args.slice(1));
+    return { command: shellArgv([...k3sExecPrefix(route), "sh", "-lc", script]), cwd: null, tty: false };
+  }
+  if (op === "argv") return { command: shellArgv([...k3sExecPrefix(route), ...args.slice(1)]), cwd: null, tty: false };
+  return { command: shellArgv([...k3sExecPrefix(route), ...args]), cwd: null, tty: false };
+}
+
+async function buildOpenPayload(argv) {
+  if (argv.length === 0) fail("schema-invalid", "tran requires a route", { help: jsonHelp() });
+  const route = parseRoute(argv[0]);
+  const parsed = route.plane === "k3s" ? await k3sCommand(route, argv.slice(1)) : await hostCommand(route, argv.slice(1));
+  return {
+    providerId: route.providerId,
+    command: parsed.command || undefined,
+    cwd: parsed.cwd || undefined,
+    tty: parsed.tty === true,
+    cols: Number(process.stdout.columns) > 0 ? Number(process.stdout.columns) : 100,
+    rows: Number(process.stdout.rows) > 0 ? Number(process.stdout.rows) : 30,
+    openTimeoutMs: Math.max(15000, Math.min(Number(process.env.UNIDESK_SSH_OPEN_TIMEOUT_MS || 60000), 60000)),
+    runtimeTimeoutMs: Math.max(1000, Math.min(Number(process.env.UNIDESK_SSH_RUNTIME_TIMEOUT_MS || process.env.UNIDESK_TRAN_RUNTIME_TIMEOUT_MS || 60000), 60000)),
+    stdinEotOnEnd: true,
+    route: route.raw,
+  };
+}
+
+async function main() {
+  const argv = process.argv.slice(2);
+  if (argv[0] === "--help" || argv[0] === "help" || argv[0] === "-h") {
+    writeJson(jsonHelp());
+    return;
+  }
+  const token = (process.env.UNIDESK_SSH_CLIENT_TOKEN || "").trim();
+  if (!token) fail("secret-unavailable", "UNIDESK_SSH_CLIENT_TOKEN is required for runner-side tran");
+  const baseUrl = baseUrlFromEnv(process.env);
+  if (!baseUrl) fail("schema-invalid", "UNIDESK_MAIN_SERVER_IP, UNIDESK_MAIN_SERVER_HOST, or UNIDESK_FRONTEND_URL is required for runner-side tran");
+  const open = await buildOpenPayload(argv);
+  await runWebSocket(open, websocketUrl(baseUrl), token);
+}
+
+async function runWebSocket(open, url, token) {
+  const ws = new WebSocket(url, { headers: { authorization: `Bearer ${token}` } });
+  let exitCode = 255;
+  let canSend = false;
+  let sessionReady = false;
+  let settled = false;
+  const pending = [];
+  const pendingInput = [];
+
+  const send = (value) => {
+    const text = JSON.stringify(value);
+    if (!canSend || ws.readyState !== WebSocket.OPEN) pending.push(text);
+    else ws.send(text);
+  };
+  const sendInput = (value) => {
+    const text = JSON.stringify(value);
+    if (!sessionReady || ws.readyState !== WebSocket.OPEN) pendingInput.push(text);
+    else ws.send(text);
+  };
+  const flush = () => {
+    while (pending.length > 0 && ws.readyState === WebSocket.OPEN) ws.send(pending.shift());
+  };
+  const flushInput = () => {
+    if (!sessionReady || ws.readyState !== WebSocket.OPEN) return;
+    while (pendingInput.length > 0) ws.send(pendingInput.shift());
+  };
+  const finish = (code) => {
+    if (settled) return;
+    settled = true;
+    clearTimeout(openTimer);
+    clearTimeout(runtimeTimer);
+    process.exit(code);
+  };
+  const openTimer = setTimeout(() => {
+    if (sessionReady || settled) return;
+    process.stderr.write("unidesk runner tran timed out waiting for provider session\n");
+    exitCode = 255;
+    try { ws.close(); } catch {}
+  }, open.openTimeoutMs);
+  const runtimeTimer = setTimeout(() => {
+    if (settled) return;
+    process.stderr.write(`UNIDESK_TRAN_TIMEOUT_HINT ${JSON.stringify({ code: "tran-runtime-timeout", level: "warning", route: open.route, timeoutMs: open.runtimeTimeoutMs, message: "tran exceeded the runtime limit; use short query plus poll semantics" })}\n`);
+    exitCode = 124;
+    try { ws.close(); } catch {}
+    finish(124);
+  }, open.runtimeTimeoutMs);
+
+  ws.addEventListener("open", () => {
+    canSend = true;
+    send({ type: "ssh.open", providerId: open.providerId, command: open.command, cwd: open.cwd, tty: open.tty, cols: open.cols, rows: open.rows });
+    flush();
+  });
+  ws.addEventListener("message", (event) => {
+    let message;
+    const text = typeof event.data === "string" ? event.data : Buffer.from(event.data).toString("utf8");
+    try {
+      message = JSON.parse(text);
+    } catch {
+      process.stderr.write(`${text}\n`);
+      return;
+    }
+    if (message.type === "ssh.dispatched") return;
+    if (message.type === "ssh.opened") {
+      sessionReady = true;
+      clearTimeout(openTimer);
+      sendInput({ type: "ssh.input", data: Buffer.from([4]).toString("base64"), encoding: "base64" });
+      sendInput({ type: "ssh.eof" });
+      flushInput();
+      return;
+    }
+    if (message.type === "ssh.data") {
+      const chunk = Buffer.from(String(message.data || ""), message.encoding === "base64" ? "base64" : "utf8");
+      if (message.stream === "stderr") process.stderr.write(chunk);
+      else process.stdout.write(chunk);
+      return;
+    }
+    if (message.type === "ssh.error") {
+      process.stderr.write(`${String(message.message || "ssh bridge error")}\n`);
+      exitCode = 255;
+      try { ws.close(); } catch {}
+      return;
+    }
+    if (message.type === "ssh.exit") {
+      exitCode = Number.isInteger(message.exitCode) ? Number(message.exitCode) : 255;
+      try { ws.close(); } catch {}
+    }
+  });
+  ws.addEventListener("close", () => finish(exitCode));
+  ws.addEventListener("error", () => {
+    process.stderr.write("unidesk runner tran websocket error\n");
+    finish(255);
+  });
+}
+
+await main();
@@ -0,0 +1,5 @@
+#!/bin/sh
+set -eu
+
+self_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
+exec "$self_dir/tran" "$@"