fix: add transactional remote patch v2

scripts/src/remote.ts

@@ -11,13 +11,16 @@ import {
   formatSshRuntimeTimeoutHint,
   formatSshRuntimeTimingHint,
   parseSshInvocation,
+  remoteCommandForRoute,
   sshFailureHint,
   sshRoutePayloadCwd,
   sshRuntimeTimeoutHint,
   sshRuntimeTimeoutMs,
   sshRuntimeTimingHint,
   wrapSshRemoteCommand,
+  type SshCaptureResult,
 } from "./ssh";
+import { runApplyPatchV2, type ApplyPatchV2Executor } from "./apply-patch-v2";
 import { codexJudgeQueryAsync, codexOutputQueryAsync, codexPrPreflightQueryAsync, codexQueuesQueryAsync, codexTaskQueryAsync, codexTasksQueryAsync, codexUnreadTriageAsync } from "./code-queue";
 import { runDecisionCenterCommandAsync } from "./decision-center";
 import {
@@ -83,6 +86,7 @@ const portOptions = new Set(["--main-server-port", "--server-port"]);
 const rootOptions = new Set(["--main-server-root", "--server-root"]);
 const keyOptions = new Set(["--main-server-key", "--server-key"]);
 const transportOptions = new Set(["--main-server-transport", "--server-transport"]);
+const remoteSshInputChunkBytes = 32 * 1024;
 
 function positivePort(raw: string, option: string): number {
   const value = Number(raw);
@@ -943,6 +947,12 @@ async function runRemoteSshWebSocket(
   const sendInput = (value: Buffer | string): void => {
     sendWhenSessionReady({ type: "ssh.input", data: Buffer.from(value).toString("base64"), encoding: "base64" });
   };
+  const sendInputChunked = (value: string): void => {
+    const buffer = Buffer.from(value, "utf8");
+    for (let offset = 0; offset < buffer.length; offset += remoteSshInputChunkBytes) {
+      sendInput(buffer.subarray(offset, Math.min(buffer.length, offset + remoteSshInputChunkBytes)));
+    }
+  };
   const sendWhenSessionReady = (value: unknown): void => {
     const text = JSON.stringify(value);
     if (!sessionReady || ws.readyState !== WebSocket.OPEN) {
@@ -1069,6 +1079,163 @@ async function runRemoteSshWebSocket(
   });
 }
 
+async function runRemoteSshWebSocketCapture(
+  session: FrontendSession,
+  invocation: ReturnType<typeof parseSshInvocation>,
+  command: string[],
+  input?: string,
+): Promise<SshCaptureResult> {
+  const remoteCommand = remoteCommandForRoute(invocation.route, command);
+  const captureInvocation = {
+    ...invocation,
+    parsed: { ...invocation.parsed, remoteCommand, requiresStdin: input !== undefined, invocationKind: "helper" as const },
+  };
+  const startedAtMs = Date.now();
+  const size = {
+    cols: Number(process.stdout.columns) > 0 ? Number(process.stdout.columns) : 100,
+    rows: Number(process.stdout.rows) > 0 ? Number(process.stdout.rows) : 30,
+  };
+  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 ws = openFrontendSshWebSocket(session);
+  let exitCode = 255;
+  let settled = false;
+  let canSend = false;
+  let sessionReady = false;
+  let stdout = "";
+  let stderr = "";
+  const pending: string[] = [];
+  const pendingSessionMessages: string[] = [];
+
+  const send = (value: unknown): void => {
+    const text = JSON.stringify(value);
+    if (!canSend || ws.readyState !== WebSocket.OPEN) {
+      pending.push(text);
+      return;
+    }
+    ws.send(text);
+  };
+  const sendWhenSessionReady = (value: unknown): void => {
+    const text = JSON.stringify(value);
+    if (!sessionReady || ws.readyState !== WebSocket.OPEN) {
+      pendingSessionMessages.push(text);
+      return;
+    }
+    ws.send(text);
+  };
+  const sendInput = (value: Buffer | string): void => {
+    sendWhenSessionReady({ type: "ssh.input", data: Buffer.from(value).toString("base64"), encoding: "base64" });
+  };
+  const flush = (): void => {
+    while (pending.length > 0 && ws.readyState === WebSocket.OPEN) ws.send(pending.shift()!);
+  };
+  const flushSessionMessages = (): void => {
+    if (!sessionReady || ws.readyState !== WebSocket.OPEN) return;
+    while (pendingSessionMessages.length > 0) ws.send(pendingSessionMessages.shift()!);
+  };
+
+  return await new Promise<SshCaptureResult>((resolve) => {
+    let killTimer: ReturnType<typeof setTimeout> | null = null;
+    const finish = (code: number): void => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(openTimer);
+      clearTimeout(runtimeTimer);
+      if (killTimer !== null) clearTimeout(killTimer);
+      const timingHint = formatSshRuntimeTimingHint(sshRuntimeTimingHint({
+        invocation: captureInvocation,
+        transport: "frontend-websocket",
+        exitCode: code,
+        startedAtMs,
+      }));
+      if (timingHint) stderr += timingHint;
+      resolve({ exitCode: code, stdout, stderr });
+    };
+    const openTimer = setTimeout(() => {
+      if (sessionReady || settled) return;
+      stderr += "unidesk remote frontend ssh bridge timed out waiting for provider session\n";
+      exitCode = 255;
+      try {
+        ws.close();
+      } catch {
+        // Ignore.
+      }
+    }, payload.openTimeoutMs);
+    const runtimeTimer = setTimeout(() => {
+      if (settled) return;
+      exitCode = 124;
+      stderr += formatSshRuntimeTimeoutHint(sshRuntimeTimeoutHint({
+        invocation: captureInvocation,
+        transport: "frontend-websocket",
+        timeoutMs: runtimeTimeoutMs,
+      }));
+      try {
+        ws.close();
+      } catch {
+        // Ignore.
+      }
+      killTimer = setTimeout(() => finish(124), 2000);
+      finish(124);
+    }, runtimeTimeoutMs);
+
+    ws.addEventListener("open", () => {
+      canSend = true;
+      send({ type: "ssh.open", ...payload });
+      flush();
+    });
+    ws.addEventListener("message", (event) => {
+      const text = webSocketDataText(event.data);
+      let message: Record<string, unknown>;
+      try {
+        message = JSON.parse(text) as Record<string, unknown>;
+      } catch {
+        stderr += `${text}\n`;
+        return;
+      }
+      if (message.type === "ssh.dispatched") return;
+      if (message.type === "ssh.opened") {
+        sessionReady = true;
+        clearTimeout(openTimer);
+        if (input !== undefined) sendInputChunked(input);
+        sendWhenSessionReady({ type: "ssh.eof" });
+        flushSessionMessages();
+        return;
+      }
+      if (message.type === "ssh.data") {
+        const chunk = Buffer.from(String(message.data ?? ""), message.encoding === "base64" ? "base64" : "utf8").toString("utf8");
+        if (message.stream === "stderr") stderr += chunk;
+        else stdout += chunk;
+        return;
+      }
+      if (message.type === "ssh.error") {
+        stderr += `${String(message.message || "ssh bridge error")}\n`;
+        exitCode = 255;
+        ws.close();
+        return;
+      }
+      if (message.type === "ssh.exit") {
+        exitCode = Number.isInteger(message.exitCode) ? Number(message.exitCode) : 255;
+        ws.close();
+      }
+    });
+    ws.addEventListener("close", () => finish(exitCode));
+    ws.addEventListener("error", () => {
+      stderr += "unidesk remote frontend ssh bridge websocket error\n";
+      finish(255);
+    });
+  });
+}
+
 export function remoteSshFrontendPlanForTest(target: string, args: string[]): Record<string, unknown> {
   const invocation = parseSshInvocation(target, args);
   return {
@@ -1086,6 +1253,12 @@ export function remoteSshFrontendPlanForTest(target: string, args: string[]): Re
 async function runRemoteSshOverFrontend(session: FrontendSession, target: string | undefined, args: string[]): Promise<number> {
   if (!target) throw new Error("remote ssh requires a route, for example: bun scripts/cli.ts --main-server-ip 74.48.78.17 ssh D601 hostname");
   const invocation = parseSshInvocation(target, args);
+  if ((args[0] ?? "") === "v2") {
+    const executor: ApplyPatchV2Executor = {
+      run: (command, input) => runRemoteSshWebSocketCapture(session, invocation, command, input),
+    };
+    return await runApplyPatchV2({ executor, stdin: process.stdin, stdout: process.stdout, argv: args.slice(1) });
+  }
   return runRemoteSshWebSocket(session, invocation);
 }