fix: stream ssh downloads over tcp-pool

scripts/src/ssh.ts

@@ -13,7 +13,12 @@ import {
   type ApplyPatchV2Executor,
   type ApplyPatchV2FileSystem,
 } from "./apply-patch-v2";
-import { isSshFileTransferOperation, runSshFileTransferOperation, type SshRemoteCommandExecutor } from "./ssh-file-transfer";
+import {
+  isSshFileTransferOperation,
+  runSshFileTransferOperation,
+  type SshRemoteCommandExecutor,
+  type SshRemoteCommandStreamHandlers,
+} from "./ssh-file-transfer";
 
 export interface ParsedSshArgs {
   remoteCommand: string | null;
@@ -2843,6 +2848,128 @@ async function runSshCaptureRemoteCommand(config: UniDeskConfig, invocation: Par
   });
 }
 
+async function runSshStreamRemoteCommand(
+  config: UniDeskConfig,
+  invocation: ParsedSshInvocation,
+  remoteCommand: string,
+  handlers: SshRemoteCommandStreamHandlers,
+  input?: string,
+): Promise<SshCaptureResult> {
+  const streamInvocation = {
+    ...invocation,
+    parsed: { ...invocation.parsed, remoteCommand, requiresStdin: input !== undefined, invocationKind: "helper" as const },
+  };
+  const startedAtMs = Date.now();
+  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 = "";
+  let streamError: unknown = null;
+  let streamWrites = Promise.resolve();
+  const queueStreamWrite = (chunk: Buffer, stream: "stdout" | "stderr"): void => {
+    streamWrites = streamWrites.then(async () => {
+      if (stream === "stdout") await handlers.onStdout(chunk);
+      else if (handlers.onStderr !== undefined) await handlers.onStderr(chunk);
+    }).catch((error) => {
+      streamError = error;
+      stderr += `unidesk ssh stream sink failed: ${error instanceof Error ? error.message : String(error)}\n`;
+      try {
+        child.kill("SIGTERM");
+      } catch {
+        // Ignore kill failures after a local stream sink error.
+      }
+    });
+  };
+  child.stdout.on("data", (chunk: Buffer) => {
+    queueStreamWrite(chunk, "stdout");
+  });
+  child.stderr.on("data", (chunk: Buffer) => {
+    stderr += chunk.toString("utf8");
+    queueStreamWrite(chunk, "stderr");
+  });
+  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);
+      void streamWrites.then(() => {
+        const finalCode = streamError === null ? exitCode : 255;
+        const timingHint = formatSshRuntimeTimingHint(sshRuntimeTimingHint({
+          invocation: streamInvocation,
+          transport: "backend-core-broker",
+          exitCode: finalCode,
+          startedAtMs,
+        }));
+        if (timingHint) stderr += timingHint;
+        resolve({ exitCode: finalCode, stdout, stderr });
+      });
+    };
+    const runtimeTimer = setTimeout(() => {
+      const hint = formatSshRuntimeTimeoutHint(sshRuntimeTimeoutHint({
+        invocation: streamInvocation,
+        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));
+  });
+}
+
 async function runRemoteSsh(config: UniDeskConfig, host: string, providerId: string, args: string[]): Promise<number> {
   const { runRemoteCli } = await import("./remote");
   return await runRemoteCli({
@@ -2950,6 +3077,7 @@ export async function runSsh(config: UniDeskConfig, providerId: string, args: st
   if (isSshFileTransferOperation(normalizedArgs)) {
     const executor: SshRemoteCommandExecutor = {
       runRemoteCommand: (remoteCommand, input) => runSshCaptureRemoteCommand(config, invocation, remoteCommand, input),
+      streamRemoteCommand: (remoteCommand, handlers, input) => runSshStreamRemoteCommand(config, invocation, remoteCommand, handlers, input),
     };
     return await runSshFileTransferOperation(invocation, normalizedArgs, executor, {
       buildRouteCommand: remoteCommandForRoute,