cadSet/cad-agent-studio/src/lib/robot-export.ts

import crypto from "node:crypto";
import { execFile } from "node:child_process";
import fs from "node:fs/promises";
import path from "node:path";
import { promisify } from "node:util";
import { artifactUrl, safeArtifactPath } from "@/lib/preview-artifacts";
import { enginePath } from "@/lib/paths";
import {
  readManifest,
  safeTaskId,
  taskDir,
  upsertManifest,
} from "@/lib/task-store";

const execFileAsync = promisify(execFile);

export type RobotExportFormat = "urdf" | "mjcf";

export type RobotExportResult = {
  status: "generated" | "cached";
  format: RobotExportFormat;
  taskId: string;
  packagePath: string;
  packageUrl: string;
  descriptionPath: string;
  descriptionUrl: string;
  meshPath: string;
  warnings: string[];
  robotStructure: "static_single_link";
};

type ExportState = {
  schema_version: "1.0";
  source_sha256: string;
  format: RobotExportFormat;
  package_path: string;
  description_path: string;
  mesh_path: string;
  warnings: string[];
};

function pythonExecutable() {
  return process.env.CAD_PYTHON
    || enginePath("text-to-cad", ".venv", "bin", "python");
}

function exporterScript() {
  return enginePath(
    "designir-pipeline",
    "scripts",
    "robot_description_export.py",
  );
}

function relativeArtifact(root: string, absolutePath: string) {
  const relative = path.relative(root, path.resolve(absolutePath));
  if (!relative || relative.startsWith("..") || path.isAbsolute(relative)) {
    throw new Error(`Robot exporter returned an artifact outside the task: ${absolutePath}`);
  }
  return safeArtifactPath(relative);
}

async function fileExists(absolutePath: string) {
  return fs.access(absolutePath).then(() => true).catch(() => false);
}

async function readCachedState(
  statePath: string,
  sourceSha256: string,
  format: RobotExportFormat,
  root: string,
) {
  try {
    const state = JSON.parse(await fs.readFile(statePath, "utf8")) as ExportState;
    if (
      state.source_sha256 !== sourceSha256
      || state.format !== format
    ) {
      return null;
    }
    const required = [
      state.package_path,
      state.description_path,
      state.mesh_path,
    ].map((value) => path.join(root, safeArtifactPath(value)));
    return (await Promise.all(required.map(fileExists))).every(Boolean)
      ? state
      : null;
  } catch {
    return null;
  }
}

function resultFromState(
  taskId: string,
  state: ExportState,
  status: "generated" | "cached",
): RobotExportResult {
  const version = state.source_sha256.slice(0, 16);
  return {
    status,
    format: state.format,
    taskId,
    packagePath: state.package_path,
    packageUrl: artifactUrl(taskId, state.package_path, version),
    descriptionPath: state.description_path,
    descriptionUrl: artifactUrl(taskId, state.description_path, version),
    meshPath: state.mesh_path,
    warnings: state.warnings,
    robotStructure: "static_single_link",
  };
}

export async function exportRobotDescription(
  requestedTaskId: string,
  format: RobotExportFormat,
): Promise<RobotExportResult> {
  const taskId = safeTaskId(requestedTaskId);
  if (format !== "urdf" && format !== "mjcf") {
    throw new Error("Robot export format must be urdf or mjcf.");
  }
  const manifest = await readManifest(taskId);
  if (!manifest) {
    throw new Error("The selected CAD task has no manifest.");
  }
  const sourcePath = safeArtifactPath(String(manifest.source?.path || ""));
  if (!sourcePath.toLowerCase().endsWith(".json")) {
    throw new Error("Robot export requires a DesignIR 3.0 JSON source.");
  }
  const root = taskDir(taskId);
  const sourceAbsolutePath = path.join(root, sourcePath);
  const sourceData = await fs.readFile(sourceAbsolutePath);
  const sourceSha256 = crypto
    .createHash("sha256")
    .update(sourceData)
    .digest("hex");
  const exportDir = path.join(root, "exports", format);
  const statePath = path.join(exportDir, "export-state.json");
  const cached = await readCachedState(
    statePath,
    sourceSha256,
    format,
    root,
  );
  if (cached) {
    return resultFromState(taskId, cached, "cached");
  }

  await fs.mkdir(exportDir, { recursive: true });
  const { stdout } = await execFileAsync(
    pythonExecutable(),
    [
      exporterScript(),
      "--designir",
      sourceAbsolutePath,
      "--output-dir",
      exportDir,
      "--format",
      format,
    ],
    {
      cwd: root,
      timeout: 900_000,
      maxBuffer: 8 * 1024 * 1024,
    },
  );
  const jsonLine = stdout
    .trim()
    .split(/\r?\n/)
    .reverse()
    .find((line) => line.trim().startsWith("{"));
  if (!jsonLine) {
    throw new Error("Robot exporter did not return a JSON result.");
  }
  const worker = JSON.parse(jsonLine) as Record<string, unknown>;
  const packagePath = relativeArtifact(root, String(worker.package || ""));
  const descriptionPath = relativeArtifact(
    root,
    String(worker.description || ""),
  );
  const meshPath = relativeArtifact(root, String(worker.mesh || ""));
  const warnings = Array.isArray(worker.warnings)
    ? worker.warnings.map(String)
    : [];
  const state: ExportState = {
    schema_version: "1.0",
    source_sha256: sourceSha256,
    format,
    package_path: packagePath,
    description_path: descriptionPath,
    mesh_path: meshPath,
    warnings,
  };
  await fs.writeFile(
    statePath,
    `${JSON.stringify(state, null, 2)}\n`,
    "utf8",
  );

  const existingArtifacts = manifest.artifacts || [];
  const nextArtifacts = existingArtifacts.filter(
    (artifact) => artifact.role !== `${format}_package`
      && artifact.role !== `${format}_description`
      && artifact.role !== `${format}_mesh`,
  );
  nextArtifacts.push(
    { path: packagePath, role: `${format}_package`, kind: "zip" },
    {
      path: descriptionPath,
      role: `${format}_description`,
      kind: format,
    },
    { path: meshPath, role: `${format}_mesh`, kind: "stl" },
  );
  await upsertManifest(taskId, {
    artifacts: nextArtifacts,
    studio: {
      ...(manifest.studio || {}),
      robotExports: {
        ...(
          manifest.studio?.robotExports
          && typeof manifest.studio.robotExports === "object"
            ? manifest.studio.robotExports
            : {}
        ),
        [format]: state,
      },
    },
  });
  return resultFromState(taskId, state, "generated");
}