fix: stream ssh downloads over tcp-pool

scripts/src/remote.ts

@@ -24,7 +24,12 @@ import {
   wrapSshRemoteCommand,
   type SshCaptureResult,
 } from "./ssh";
-import { isSshFileTransferOperation, runSshFileTransferOperation, type SshRemoteCommandExecutor } from "./ssh-file-transfer";
+import {
+  isSshFileTransferOperation,
+  runSshFileTransferOperation,
+  type SshRemoteCommandExecutor,
+  type SshRemoteCommandStreamHandlers,
+} from "./ssh-file-transfer";
 import { runApplyPatchV2, type ApplyPatchV2Executor } from "./apply-patch-v2";
 import { codexJudgeQueryAsync, codexOutputQueryAsync, codexPrPreflightQueryAsync, codexQueuesQueryAsync, codexTaskQueryAsync, codexTasksQueryAsync, codexUnreadTriageAsync } from "./code-queue";
 import { runDecisionCenterCommandAsync } from "./decision-center";
@@ -1194,6 +1199,12 @@ async function runRemoteSshWebSocketCaptureRemoteCommand(
   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 flush = (): void => {
     while (pending.length > 0 && ws.readyState === WebSocket.OPEN) ws.send(pending.shift()!);
   };
@@ -1294,6 +1305,193 @@ async function runRemoteSshWebSocketCaptureRemoteCommand(
   });
 }
 
+async function runRemoteSshWebSocketStreamRemoteCommand(
+  session: FrontendSession,
+  invocation: ReturnType<typeof parseSshInvocation>,
+  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 = {
+    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 = "";
+  let streamError: unknown = null;
+  let streamWrites = Promise.resolve();
+  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 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 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()!);
+  };
+  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 remote frontend ssh stream sink failed: ${error instanceof Error ? error.message : String(error)}\n`;
+      exitCode = 255;
+      try {
+        ws.close();
+      } catch {
+        // Ignore close failures after the local stream sink has failed.
+      }
+    });
+  };
+
+  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);
+      void streamWrites.then(() => {
+        const finalCode = streamError === null ? code : 255;
+        const timingHint = formatSshRuntimeTimingHint(sshRuntimeTimingHint({
+          invocation: streamInvocation,
+          transport: "frontend-websocket",
+          exitCode: finalCode,
+          startedAtMs,
+        }));
+        if (timingHint) stderr += timingHint;
+        resolve({ exitCode: finalCode, 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: streamInvocation,
+        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");
+        if (message.stream === "stderr") {
+          stderr += chunk.toString("utf8");
+          queueStreamWrite(chunk, "stderr");
+        } else {
+          queueStreamWrite(chunk, "stdout");
+        }
+        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 {
@@ -1316,6 +1514,7 @@ async function runRemoteSshOverFrontend(session: FrontendSession, target: string
   if (isSshFileTransferOperation(normalizedArgs)) {
     const executor: SshRemoteCommandExecutor = {
       runRemoteCommand: (remoteCommand, input) => runRemoteSshWebSocketCaptureRemoteCommand(session, invocation, remoteCommand, input),
+      streamRemoteCommand: (remoteCommand, handlers, input) => runRemoteSshWebSocketStreamRemoteCommand(session, invocation, remoteCommand, handlers, input),
     };
     return await runSshFileTransferOperation(invocation, normalizedArgs, executor, {
       buildRouteCommand: remoteCommandForRoute,