pikasTech-unidesk/scripts/src/agentrun/target-execution.ts

import { readFileSync } from "node:fs";
import type { UniDeskConfig } from "../config";
import type { RenderedCliResult } from "../output";
import { resolveAgentRunLaneTarget } from "../agentrun-lanes";
import { runSshCommandCapture } from "../ssh";
import { readAgentRunClientConfig } from "./config";
import { renderMachine, renderedCliResult } from "./render";
import type { AgentRunResourceOptions } from "./utils";

export async function runAgentRunResourceThroughTarget(
  config: UniDeskConfig | null,
  command: string,
  canonicalArgs: string[],
  options: AgentRunResourceOptions,
): Promise<RenderedCliResult | null> {
  const client = readAgentRunClientConfig();
  if (client.client.transport !== "target-trans") return null;
  if (config === null) return null;
  const execution = client.client.targetExecution;
  if (execution === null) throw new Error(`${client.sourcePath}: client.targetExecution is required for target-trans`);
  const { configPath, spec } = resolveAgentRunLaneTarget({ node: options.node, lane: options.lane });
  const activeTarget = process.env[execution.markerEnv] ?? null;
  if (activeTarget === spec.nodeId) return null;
  if (activeTarget !== null) {
    throw new Error(`${execution.markerEnv}=${activeTarget} cannot re-enter a second Target ${spec.nodeId}`);
  }
  const route = `${spec.nodeRoute}:${spec.nodeUnideskWorkspace}`;
  const targetArgs = [...canonicalArgs];
  if (!hasOption(targetArgs, "target")) targetArgs.push("--target", spec.nodeId);
  if (!hasOption(targetArgs, "lane")) targetArgs.push("--lane", spec.lane);
  const stdin = targetExecutionStdin(targetArgs);
  const capture = await runSshCommandCapture(config, route, [
    "argv",
    "env",
    `${execution.markerEnv}=${spec.nodeId}`,
    execution.executable,
    execution.cliPath,
    "agentrun",
    ...targetArgs,
  ], stdin);
  const disclosure = {
    transport: "target-trans",
    target: spec.nodeId,
    lane: spec.lane,
    route,
    configPath,
    unideskWorkspace: spec.nodeUnideskWorkspace,
    markerEnv: execution.markerEnv,
    mutation: false,
    valuesPrinted: false,
  };
  if (capture.exitCode !== 0) {
    const detail = boundedRemoteFailure(capture.stderr || capture.stdout);
    return renderedCliResult(
      false,
      command,
      [
        `ERROR  target-trans/${spec.nodeId}`,
        `route: ${route}`,
        `exitCode: ${capture.exitCode}`,
        `message: ${detail || "Target CLI did not return output"}`,
      ].join("\n"),
    );
  }

  const output = capture.stdout.trimEnd();
  if (options.raw || options.output === "json") {
    const payload = parseTargetMachineOutput(output, "json");
    return renderMachine(command, attachTargetExecution(payload, disclosure), "json", machineOutputOk(payload));
  }
  if (options.output === "yaml") {
    const payload = parseTargetMachineOutput(output, "yaml");
    return renderMachine(command, attachTargetExecution(payload, disclosure), "yaml", machineOutputOk(payload));
  }
  return renderedCliResult(
    true,
    command,
    [`TargetExecution: trans ${route} lane=${spec.lane}`, output].filter(Boolean).join("\n"),
  );
}

function hasOption(args: string[], name: string): boolean {
  const flag = `--${name}`;
  return args.some((arg) => arg === flag || arg.startsWith(`${flag}=`));
}

function targetExecutionStdin(args: string[]): string | undefined {
  const stdinFlag = args.some((arg) => arg === "--prompt-stdin" || arg === "--stdin"
    || arg === "--json-stdin" || arg === "--yaml-stdin");
  const stdinFile = optionIsStdin(args, "prompt-file") || optionIsStdin(args, "json-file")
    || optionIsStdin(args, "yaml-file") || optionIsStdin(args, "file") || optionIsStdin(args, "filename")
    || args.some((arg, index) => arg === "-f" && args[index + 1] === "-");
  return stdinFlag || stdinFile ? readFileSync(0, "utf8") : undefined;
}

function optionIsStdin(args: string[], name: string): boolean {
  const flag = `--${name}`;
  return args.some((arg, index) => (arg === flag && args[index + 1] === "-") || arg === `${flag}=-`);
}

function parseTargetMachineOutput(output: string, mode: "json" | "yaml"): unknown {
  if (output.length === 0) throw new Error(`Target CLI returned empty ${mode} output`);
  try {
    return mode === "json" ? JSON.parse(output) : Bun.YAML.parse(output);
  } catch (error) {
    throw new Error(`Target CLI returned invalid ${mode} output: ${error instanceof Error ? error.message : String(error)}`);
  }
}

function attachTargetExecution(value: unknown, disclosure: Record<string, unknown>): Record<string, unknown> {
  if (typeof value === "object" && value !== null && !Array.isArray(value)) {
    return { ...(value as Record<string, unknown>), targetExecution: disclosure };
  }
  return { data: value, targetExecution: disclosure };
}

function machineOutputOk(value: unknown): boolean {
  return !(typeof value === "object" && value !== null && (value as Record<string, unknown>).ok === false);
}

function boundedRemoteFailure(value: string): string {
  return value
    .replace(/https?:\/\/[^@\s]+@/gu, "https://[redacted]@")
    .replace(/(token|password|authorization)[=:][^\s]+/giu, "$1=[redacted]")
    .replace(/\s+/gu, " ")
    .trim()
    .slice(0, 600);
}