paperclip/packages/adapter-utils/src/command-managed-runtime.ts

import path from "node:path";
import {
  prepareSandboxManagedRuntime,
  type PreparedSandboxManagedRuntime,
  type SandboxManagedRuntimeAsset,
  type SandboxManagedRuntimeClient,
  type SandboxRemoteExecutionSpec,
} from "./sandbox-managed-runtime.js";
import { preferredShellForSandbox, shellCommandArgs } from "./sandbox-shell.js";
import type { RunProcessResult } from "./server-utils.js";
import type { RuntimeProgressSink } from "./runtime-progress.js";

export interface CommandManagedRuntimeRunner {
  /**
   * True only when `execute({ stdin })` can surface useful in-flight progress
   * for a single stdin-backed command. Provider-backed sandbox runners usually
   * complete the entire RPC before returning, so they should leave this false
   * and let the caller choose a chunked upload path when progress is requested.
   */
  supportsSingleStreamStdinProgress?: boolean;
  execute(input: {
    command: string;
    args?: string[];
    cwd?: string;
    env?: Record<string, string>;
    stdin?: string;
    timeoutMs?: number;
    onLog?: (stream: "stdout" | "stderr", chunk: string) => Promise<void>;
    onSpawn?: (meta: { pid: number; startedAt: string }) => Promise<void>;
  }): Promise<RunProcessResult>;
}

export interface CommandManagedRuntimeSpec {
  providerKey?: string | null;
  shellCommand?: "bash" | "sh" | null;
  leaseId?: string | null;
  remoteCwd: string;
  timeoutMs?: number | null;
}

export type CommandManagedRuntimeAsset = SandboxManagedRuntimeAsset;

function shellQuote(value: string) {
  return `'${value.replace(/'/g, `'"'"'`)}'`;
}

function mergeRuntimeExcludes(entries: string[] | undefined): string[] {
  return [...new Set([".paperclip-runtime", ...(entries ?? [])])];
}

// Largest base64 body we hand to the runner as a single stdin string. Normal
// multi-MB workspace/asset tarballs stay well under this and upload in one
// round-trip; anything larger uses the bounded chunked-append fallback so a
// runaway stdin string can't blow the runner/provider RPC limits.
const REMOTE_WRITE_SINGLE_STREAM_MAX_BASE64_BYTES = 96 * 1024 * 1024;
// Fallback chunk size (base64 bytes). Kept a multiple of 4 so each chunk is a
// self-contained base64 unit that decodes cleanly on its own.
const REMOTE_WRITE_FALLBACK_BASE64_CHUNK_SIZE = 4 * 1024 * 1024;

function toBuffer(bytes: Buffer | Uint8Array | ArrayBuffer): Buffer {
  if (Buffer.isBuffer(bytes)) return bytes;
  if (bytes instanceof ArrayBuffer) return Buffer.from(bytes);
  return Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength);
}

function requireSuccessfulResult(result: RunProcessResult, action: string): void {
  if (result.exitCode === 0 && !result.timedOut) return;
  const stderr = result.stderr.trim();
  const detail = stderr.length > 0 ? `: ${stderr}` : "";
  throw new Error(`${action} failed with exit code ${result.exitCode ?? "null"}${detail}`);
}

export function createCommandManagedRuntimeClient(input: {
  runner: CommandManagedRuntimeRunner;
  commandCwd: string;
  timeoutMs: number;
  shellCommand?: "bash" | "sh" | null;
}): SandboxManagedRuntimeClient {
  const shellCommand = preferredShellForSandbox(input.shellCommand);
  const runShell = async (
    script: string,
    opts: {
      stdin?: string;
      timeoutMs?: number;
      onLog?: (stream: "stdout" | "stderr", chunk: string) => Promise<void>;
    } = {},
  ) => {
    const result = await input.runner.execute({
      command: shellCommand,
      args: shellCommandArgs(script),
      cwd: input.commandCwd,
      stdin: opts.stdin,
      timeoutMs: opts.timeoutMs ?? input.timeoutMs,
      onLog: opts.onLog,
    });
    requireSuccessfulResult(result, script);
    return result;
  };

  return {
    makeDir: async (remotePath) => {
      await runShell(`mkdir -p ${shellQuote(remotePath)}`);
    },
    writeFile: async (remotePath, bytes, options) => {
      const buffer = toBuffer(bytes);
      const total = buffer.byteLength;
      const body = buffer.toString("base64");
      const remoteDir = path.posix.dirname(remotePath);
      const remoteTempPath = `${remotePath}.paperclip-upload`;
      const canUseSingleStreamProgressPath =
        !options?.onProgress || input.runner.supportsSingleStreamStdinProgress === true;

      // Primary path: a single round-trip. Stream the entire base64 body to one
      // `base64 -d` process via stdin, decode straight into a temp file, then
      // atomically rename into place. This replaces the previous loop that did
      // one `printf >> tmpfile` shell round-trip per 32 KB — thousands of serial
      // processes for a large workspace — with exactly one process.
      if (
        body.length <= REMOTE_WRITE_SINGLE_STREAM_MAX_BASE64_BYTES &&
        canUseSingleStreamProgressPath
      ) {
        await options?.onProgress?.(0, total);
        await runShell(
          `mkdir -p ${shellQuote(remoteDir)} && ` +
            `base64 -d > ${shellQuote(remoteTempPath)} && ` +
            `mv -f ${shellQuote(remoteTempPath)} ${shellQuote(remotePath)}`,
          { stdin: body },
        );
        await options?.onProgress?.(total, total);
        return;
      }

      // Bounded fallback for payloads too large to hand the runner as one stdin
      // string: append the base64 body to a remote temp file in large chunks
      // (orders of magnitude fewer round-trips than the old 32 KB loop), decoding
      // each self-contained chunk on arrival and emitting progress per write,
      // then atomically rename into place.
      await runShell(
        `mkdir -p ${shellQuote(remoteDir)} && ` +
          `rm -f ${shellQuote(remoteTempPath)} && : > ${shellQuote(remoteTempPath)}`,
      );
      for (let offset = 0; offset < body.length; offset += REMOTE_WRITE_FALLBACK_BASE64_CHUNK_SIZE) {
        const chunk = body.slice(offset, offset + REMOTE_WRITE_FALLBACK_BASE64_CHUNK_SIZE);
        await runShell(`base64 -d >> ${shellQuote(remoteTempPath)}`, { stdin: chunk });
        const decodedSoFar = Math.min(total, Math.floor(((offset + chunk.length) * 3) / 4));
        await options?.onProgress?.(decodedSoFar, total);
      }
      await runShell(`mv -f ${shellQuote(remoteTempPath)} ${shellQuote(remotePath)}`);
      await options?.onProgress?.(total, total);
    },
    readFile: async (remotePath, options) => {
      // Decoded file size up front so download progress can be a percentage.
      // Only paid when a progress hook is attached.
      let totalBytes: number | null = null;
      if (options?.onProgress) {
        const sizeResult = await runShell(`wc -c < ${shellQuote(remotePath)}`);
        const parsed = Number.parseInt(sizeResult.stdout.trim(), 10);
        totalBytes = Number.isFinite(parsed) && parsed >= 0 ? parsed : null;
      }

      // Stream the remote `base64` stdout through a byte counter, decoding each
      // 4-char-aligned slice incrementally rather than buffering the whole
      // response as one string. Falls back to the buffered result when the
      // runner does not surface incremental stdout.
      const decodedChunks: Buffer[] = [];
      let b64Remainder = "";
      let decodedSoFar = 0;
      let streamed = false;
      const result = await runShell(`base64 < ${shellQuote(remotePath)}`, {
        onLog: async (stream, chunk) => {
          if (stream !== "stdout") return;
          streamed = true;
          const data = b64Remainder + chunk.replace(/\s+/g, "");
          const alignedLen = data.length - (data.length % 4);
          if (alignedLen > 0) {
            const decoded = Buffer.from(data.slice(0, alignedLen), "base64");
            decodedChunks.push(decoded);
            decodedSoFar += decoded.byteLength;
          }
          b64Remainder = data.slice(alignedLen);
          if (options?.onProgress) {
            const done = totalBytes != null ? Math.min(totalBytes, decodedSoFar) : decodedSoFar;
            await options.onProgress(done, totalBytes);
          }
        },
      });

      let out: Buffer;
      if (streamed) {
        if (b64Remainder.length > 0) {
          decodedChunks.push(Buffer.from(b64Remainder, "base64"));
        }
        out = Buffer.concat(decodedChunks);
      } else {
        out = Buffer.from(result.stdout.replace(/\s+/g, ""), "base64");
      }
      await options?.onProgress?.(out.byteLength, totalBytes ?? out.byteLength);
      return out;
    },
    listFiles: async (remotePath) => {
      const result = await runShell(
        `if [ -d ${shellQuote(remotePath)} ]; then ` +
          `for entry in ${shellQuote(remotePath)}/*; do ` +
          `[ -f "$entry" ] || continue; ` +
          `basename "$entry"; ` +
          `done; ` +
        `fi`,
      );
      return result.stdout
        .split(/\r?\n/)
        .map((entry) => entry.trim())
        .filter((entry) => entry.length > 0)
        .sort((left, right) => left.localeCompare(right));
    },
    remove: async (remotePath) => {
      const result = await input.runner.execute({
        command: shellCommand,
        args: shellCommandArgs(`rm -rf ${shellQuote(remotePath)}`),
        cwd: input.commandCwd,
        timeoutMs: input.timeoutMs,
      });
      requireSuccessfulResult(result, `remove ${remotePath}`);
    },
    run: async (command, options) => {
      const result = await input.runner.execute({
        command: shellCommand,
        args: shellCommandArgs(command),
        cwd: input.commandCwd,
        timeoutMs: options.timeoutMs,
      });
      requireSuccessfulResult(result, command);
    },
  };
}

export async function prepareCommandManagedRuntime(input: {
  runner: CommandManagedRuntimeRunner;
  spec: CommandManagedRuntimeSpec;
  adapterKey: string;
  workspaceLocalDir: string;
  workspaceRemoteDir?: string;
  workspaceExclude?: string[];
  preserveAbsentOnRestore?: string[];
  assets?: CommandManagedRuntimeAsset[];
  installCommand?: string | null;
  /** When provided alongside `installCommand`, skip the install if `command -v <detectCommand>` succeeds. */
  detectCommand?: string | null;
  // Upload progress sink. Forwarded to prepareSandboxManagedRuntime; the child
  // task wires it into the byte-counting writeFile/readFile transport.
  onProgress?: RuntimeProgressSink;
}): Promise<PreparedSandboxManagedRuntime> {
  const timeoutMs = input.spec.timeoutMs && input.spec.timeoutMs > 0 ? input.spec.timeoutMs : 300_000;
  const workspaceRemoteDir = input.workspaceRemoteDir ?? input.spec.remoteCwd;
  // Managed-runtime sync/restore scripts use absolute paths throughout, so
  // run them from a stable cwd. The target workspace itself may be removed or
  // recreated during a run, which breaks shell startup if we chdir into it.
  const commandCwd = "/";
  const runtimeSpec: SandboxRemoteExecutionSpec = {
    transport: "sandbox",
    provider: input.spec.providerKey ?? "sandbox",
    sandboxId: input.spec.leaseId ?? "managed",
    remoteCwd: workspaceRemoteDir,
    timeoutMs,
    apiKey: null,
  };
  const client = createCommandManagedRuntimeClient({
    runner: input.runner,
    commandCwd,
    timeoutMs,
    shellCommand: input.spec.shellCommand,
  });
  const shellCommand = preferredShellForSandbox(input.spec.shellCommand);

  if (input.installCommand?.trim()) {
    const installCommand = input.installCommand.trim();
    const detectCommand = input.detectCommand?.trim();
    // Skip the install when the binary is already on PATH. Without this
    // probe the install runs unconditionally on every execute() call (and
    // also runs a second time after `ensureAdapterExecutionTargetCommandResolvable`
    // has already installed it during the resolvability gate).
    if (detectCommand) {
      const probe = await input.runner.execute({
        command: shellCommand,
        args: shellCommandArgs(`command -v ${shellQuote(detectCommand)} >/dev/null 2>&1`),
        cwd: commandCwd,
        timeoutMs,
      });
      if (!probe.timedOut && (probe.exitCode ?? 1) === 0) {
        return await prepareSandboxManagedRuntime({
          spec: runtimeSpec,
          client,
          adapterKey: input.adapterKey,
          workspaceLocalDir: input.workspaceLocalDir,
          workspaceRemoteDir,
          workspaceExclude: mergeRuntimeExcludes(input.workspaceExclude),
          preserveAbsentOnRestore: input.preserveAbsentOnRestore,
          assets: input.assets,
          onProgress: input.onProgress,
        });
      }
    }
    const result = await input.runner.execute({
      command: shellCommand,
      args: shellCommandArgs(installCommand),
      cwd: commandCwd,
      timeoutMs,
    });
    // A failed install is not always fatal: the CLI may already be on PATH
    // from a previous lease, the template image, or another path entry. Log
    // and continue rather than aborting the agent run; downstream code that
    // exec's the CLI will surface a clear "command not found" if it is in
    // fact missing. The test path's `maybeRunSandboxInstallCommand` already
    // honors this contract — keep them consistent.
    if (result.timedOut || (result.exitCode ?? 0) !== 0) {
      const tail = (text: string) =>
        text.split(/\r?\n/).filter((line) => line.trim().length > 0).slice(-3).join(" | ").slice(0, 480);
      const reason = result.timedOut ? "timed out" : `exited ${result.exitCode ?? "?"}`;
      console.warn(
        `[paperclip] managed-runtime install command ${reason}: ${installCommand} :: ${tail(result.stderr || result.stdout)}`,
      );
    }
  }

  return await prepareSandboxManagedRuntime({
    spec: runtimeSpec,
    client,
    adapterKey: input.adapterKey,
    workspaceLocalDir: input.workspaceLocalDir,
    workspaceRemoteDir,
    workspaceExclude: mergeRuntimeExcludes(input.workspaceExclude),
    preserveAbsentOnRestore: input.preserveAbsentOnRestore,
    assets: input.assets,
    onProgress: input.onProgress,
  });
}