fix: add transactional remote patch v2

2026-05-26 17:38:09 +00:00
parent 6a596bc452
commit 98715ce2be
8 changed files with 1278 additions and 40 deletions
@@ -1,5 +1,6 @@
 import { spawn } from "node:child_process";
 import { type UniDeskConfig, repoRoot } from "./config";
+import { isApplyPatchV2HelpArgs, runApplyPatchV2, type ApplyPatchV2Executor } from "./apply-patch-v2";

 export interface ParsedSshArgs {
  remoteCommand: string | null;
@@ -30,6 +31,12 @@ export interface ParsedSshInvocation {
  parsed: ParsedSshArgs;
 }

+export interface SshCaptureResult {
+  exitCode: number;
+  stdout: string;
+  stderr: string;
+}
+
 export interface SshFailureHint {
  code: "ssh-like-command-friction";
  providerId: string;
@@ -86,6 +93,7 @@ const legacyK3sOperationRouteSegments = new Set([
  "exec",
  "script",
  "apply-patch",
+  "v2",
  "logs",
  "get",
  "describe",
@@ -821,6 +829,12 @@ export function parseSshArgs(args: string[]): ParsedSshArgs {
    const toolArgs = ["apply_patch", ...args.slice(1)];
    return { remoteCommand: shellArgv(toolArgs), requiresStdin: true, invocationKind: "helper", requiredHelpers: ["apply_patch"] };
  }
+  if (subcommand === "v2") {
+    if (isApplyPatchV2HelpArgs(args.slice(1))) {
+      return { remoteCommand: null, requiresStdin: false, invocationKind: "helper" };
+    }
+    return { remoteCommand: null, requiresStdin: true, invocationKind: "helper" };
+  }
  if (subcommand === "py") {
    return { remoteCommand: buildPythonStdinCommand(args.slice(1)), requiresStdin: true, invocationKind: "helper" };
  }
@@ -1153,11 +1167,11 @@ function k3sOperationInRouteMessage(target: string, operation: string): string {
  return `ssh k3s route must locate a target only; put operation "${operation}" after the route, for example "ssh ${providerId}:k3s ${operationExample}" or "ssh ${providerId}:k3s:<namespace>:<workload> ${operationExample}" instead of "${target}"`;
 }

-function shellArgv(args: string[]): string {
+export function shellArgv(args: string[]): string {
  return args.map(shellQuote).join(" ");
 }

-function shellQuote(value: string): string {
+export function shellQuote(value: string): string {
  return `'${value.replace(/'/g, `'\\''`)}'`;
 }

@@ -1272,6 +1286,9 @@ function parseK3sControlPlaneOperation(route: ParsedSshRoute, args: string[]): P
  if (operation === "apply-patch" || operation === "patch") {
    throw new Error(`ssh ${route.providerId}:k3s apply-patch requires a workload route: ssh ${route.providerId}:k3s:<namespace>:<workload> apply-patch`);
  }
+  if (operation === "v2") {
+    throw new Error(`ssh ${route.providerId}:k3s v2 requires a workload route: ssh ${route.providerId}:k3s:<namespace>:<workload> v2`);
+  }
  if (operation === "script" || operation === "sh") {
    return { remoteCommand: buildK3sScriptCommand(args.slice(1)), requiresStdin: true, invocationKind: "helper" };
  }
@@ -1298,6 +1315,7 @@ function parseK3sTargetOperation(route: ParsedSshRoute, args: string[]): ParsedS
  const operation = args[0] ?? "";
  const operationArgs = args.slice(1);
  if (operation === "apply-patch" || operation === "patch") return buildK3sApplyPatchCommand([...targetArgs, ...operationArgs]);
+  if (operation === "v2") return { remoteCommand: null, requiresStdin: true, invocationKind: "helper" };
  if (operation === "script") return { remoteCommand: buildK3sScriptCommand([...targetArgs, ...operationArgs]), requiresStdin: true, invocationKind: "helper" };
  if (operation === "shell") {
    const parsed = parseShellStringOperationArgs(operationArgs, `ssh ${route.raw} shell`);
@@ -1326,6 +1344,10 @@ function buildK3sTargetObjectCommand(action: "get" | "describe", route: ParsedSs
  return shellArgv(["env", `KUBECONFIG=${nativeK3sKubeconfig}`, "kubectl", action, "-n", route.namespace, normalizeK3sRouteResource(route.resource), ...args]);
 }

+function buildK3sTargetCommand(route: ParsedSshRoute, command: string[]): string {
+  return buildK3sExecCommand([...k3sRouteTargetArgs(route), "--", ...command]);
+}
+
 function k3sRouteTargetArgs(route: ParsedSshRoute): string[] {
  if (route.namespace === null) throw new Error(`ssh route ${route.raw} requires a namespace segment`);
  if (route.resource === null) throw new Error(`ssh route ${route.raw} requires a workload or pod segment`);
@@ -1701,6 +1723,7 @@ function buildShellCommand(args: string[]): ParsedSshArgs {
  }
  if (directArgvMode) {
    if (scriptArgs.length === 0) throw new Error("ssh script -- requires a command");
+    if (scriptArgs.length === 1) return { remoteCommand: shellArgv([shell, "-c", scriptArgs[0] ?? ""]), requiresStdin: false, invocationKind: "helper" };
    return { remoteCommand: shellArgv(scriptArgs), requiresStdin: false, invocationKind: "argv" };
  }
  return { remoteCommand: shellArgv([shell, "-s", "--", ...scriptArgs]), requiresStdin: true, invocationKind: "helper" };
@@ -2032,9 +2055,127 @@ function terminalSize(): { cols: number; rows: number } {
  };
 }

+export function remoteCommandForRoute(route: ParsedSshRoute, command: string[]): string {
+  if (route.plane === "k3s") return buildK3sTargetCommand(route, command);
+  if (route.plane === "win") throw new Error(`ssh v2 does not support win routes yet: ${route.raw}`);
+  return shellArgv(command);
+}
+
+async function runSshCaptureCommand(config: UniDeskConfig, invocation: ParsedSshInvocation, command: string[], input?: string): Promise<SshCaptureResult> {
+  const startedAtMs = Date.now();
+  const remoteCommand = remoteCommandForRoute(invocation.route, command);
+  const size = terminalSize();
+  const runtimeTimeoutMs = sshRuntimeTimeoutMs();
+  const payload = {
+    providerId: invocation.providerId,
+    command: wrapSshRemoteCommand(remoteCommand),
+    cwd: sshRoutePayloadCwd(invocation.route),
+    tty: false,
+    stdinEotOnEnd: false,
+    openTimeoutMs: Math.max(15000, Number(process.env.UNIDESK_SSH_OPEN_TIMEOUT_MS || 60000)),
+    runtimeTimeoutMs,
+    cols: size.cols,
+    rows: size.rows,
+  };
+  const payloadJson = JSON.stringify(payload);
+  const encodedBrokerSource = Buffer.from(brokerSource(), "utf8").toString("base64");
+  const script = [
+    "set -eu",
+    `payload=${shellQuote(payloadJson)}`,
+    "if command -v backend-core >/dev/null 2>&1; then",
+    '  exec backend-core --ssh-broker "$payload"',
+    "fi",
+    `export UNIDESK_SSH_BROKER_URL=${shellQuote("ws://127.0.0.1:8080/ws/ssh")}`,
+    'broker_js="$(mktemp "${TMPDIR:-/tmp}/unidesk-ssh-broker.XXXXXX")"',
+    'trap \'rm -f "$broker_js"\' EXIT',
+    `printf %s ${shellQuote(encodedBrokerSource)} | base64 -d >"$broker_js"`,
+    'bun "$broker_js" "$payload"',
+  ].join("\n");
+  const child = spawn("docker", ["exec", "-i", "unidesk-backend-core", "sh", "-c", script], {
+    cwd: repoRoot,
+    stdio: ["pipe", "pipe", "pipe"],
+  });
+  if (input !== undefined) {
+    writeChunkedStdin(child.stdin, input);
+  } else {
+    child.stdin.end();
+  }
+  let stdout = "";
+  let stderr = "";
+  child.stdout.on("data", (chunk: Buffer) => {
+    stdout += chunk.toString("utf8");
+  });
+  child.stderr.on("data", (chunk: Buffer) => {
+    stderr += chunk.toString("utf8");
+  });
+  return await new Promise<SshCaptureResult>((resolve) => {
+    let settled = false;
+    let killTimer: NodeJS.Timeout | null = null;
+    const finish = (exitCode: number): void => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(runtimeTimer);
+      if (killTimer !== null) clearTimeout(killTimer);
+      const timingHint = formatSshRuntimeTimingHint(sshRuntimeTimingHint({
+        invocation,
+        transport: "backend-core-broker",
+        exitCode,
+        startedAtMs,
+      }));
+      if (timingHint) stderr += timingHint;
+      resolve({ exitCode, stdout, stderr });
+    };
+    const runtimeTimer = setTimeout(() => {
+      const hint = formatSshRuntimeTimeoutHint(sshRuntimeTimeoutHint({
+        invocation,
+        transport: "backend-core-broker",
+        timeoutMs: runtimeTimeoutMs,
+      }));
+      stderr += hint;
+      try {
+        child.kill("SIGTERM");
+      } catch {
+        // Ignore.
+      }
+      killTimer = setTimeout(() => {
+        try {
+          child.kill("SIGKILL");
+        } catch {
+          // Ignore.
+        }
+      }, 2000);
+      finish(124);
+    }, runtimeTimeoutMs);
+    child.on("error", (error) => {
+      stderr += `unidesk ssh failed to start broker: ${error.message}\n`;
+      finish(255);
+    });
+    child.on("close", (code) => finish(code ?? 255));
+  });
+}
+
+function writeChunkedStdin(stdin: NodeJS.WritableStream, input: string): void {
+  const buffer = Buffer.from(input, "utf8");
+  const chunkSize = 32 * 1024;
+  for (let offset = 0; offset < buffer.length; offset += chunkSize) {
+    stdin.write(buffer.subarray(offset, Math.min(buffer.length, offset + chunkSize)));
+  }
+  stdin.end();
+}
+
 export async function runSsh(config: UniDeskConfig, providerId: string, args: string[]): Promise<number> {
  const invocation = parseSshInvocation(providerId, args);
  const parsed = invocation.parsed;
+  const operationName = args[0] ?? "";
+  if (operationName === "v2") {
+    const executor: ApplyPatchV2Executor = { run: (command, input) => runSshCaptureCommand(config, invocation, command, input) };
+    return await runApplyPatchV2({
+      executor,
+      stdin: process.stdin,
+      stdout: process.stdout,
+      argv: args.slice(1),
+    });
+  }
  const startedAtMs = Date.now();
  const size = terminalSize();
  const openTimeoutMs = Math.max(15000, Number(process.env.UNIDESK_SSH_OPEN_TIMEOUT_MS || 60000));
@@ -2059,8 +2200,10 @@ export async function runSsh(config: UniDeskConfig, providerId: string, args: st
    '  exec backend-core --ssh-broker "$payload"',
    "fi",
    `export UNIDESK_SSH_BROKER_URL=${shellQuote("ws://127.0.0.1:8080/ws/ssh")}`,
-    `printf %s ${shellQuote(encodedBrokerSource)} | base64 -d >/tmp/unidesk-ssh-broker.js`,
-    "exec bun /tmp/unidesk-ssh-broker.js \"$payload\"",
+    'broker_js="$(mktemp "${TMPDIR:-/tmp}/unidesk-ssh-broker.XXXXXX")"',
+    'trap \'rm -f "$broker_js"\' EXIT',
+    `printf %s ${shellQuote(encodedBrokerSource)} | base64 -d >"$broker_js"`,
+    'bun "$broker_js" "$payload"',
  ].join("\n");
  const child = spawn("docker", [
    "exec",