agent_core/tools/cadquery_tool.py

import json
import subprocess
import sys
import textwrap
from pathlib import Path

from agent_core.config import DEFAULT_CADQUERY_OUTPUT_PATH, WORKDIR
from agent_core.outputs import resolve_run_output, update_latest_link, workspace_relative, write_manifest
from agent_core.utils import json_response


CADQUERY_RUNNER = r"""
import contextlib
import io
import json
import sys
import traceback
from pathlib import Path

stdout_buffer = io.StringIO()
stderr_buffer = io.StringIO()


def tail(text, limit=4000):
    if not text:
        return ""
    return text[-limit:]


def fail(stage, exc=None, error_type=None, error=None):
    payload = {
        "ok": False,
        "stage": stage,
        "error_type": error_type or (type(exc).__name__ if exc else "Error"),
        "error": error or (str(exc) if exc else ""),
        "traceback_tail": tail(traceback.format_exc() if exc else ""),
        "stdout": tail(stdout_buffer.getvalue()),
        "stderr": tail(stderr_buffer.getvalue()),
    }
    print(json.dumps(payload, ensure_ascii=False))


try:
    request = json.loads(sys.stdin.read())
    code = request["code"]
    output_path = Path(request["output_path"])
    preview_path = Path(request["preview_path"])
except Exception as exc:
    fail("input", exc)
    sys.exit(0)

try:
    import cadquery as cq
except Exception as exc:
    fail("import", exc)
    sys.exit(0)

try:
    compiled = compile(code, "<cadquery_code>", "exec")
except Exception as exc:
    fail("syntax", exc)
    sys.exit(0)

namespace = {
    "__name__": "__cadquery_user_code__",
    "cq": cq,
}

try:
    with contextlib.redirect_stdout(stdout_buffer), contextlib.redirect_stderr(stderr_buffer):
        exec(compiled, namespace)
except Exception as exc:
    fail("execution", exc)
    sys.exit(0)

if "result" not in namespace:
    fail("result", error_type="MissingResult", error="CadQuery code must assign the final model to variable 'result'.")
    sys.exit(0)

result = namespace["result"]
if result is None:
    fail("result", error_type="InvalidResult", error="'result' is None.")
    sys.exit(0)

exportable_types = tuple(
    value
    for name in ("Workplane", "Shape", "Assembly", "Sketch")
    for value in [getattr(cq, name, None)]
    if isinstance(value, type)
)

is_exportable = isinstance(result, exportable_types)
if isinstance(result, (list, tuple)) and result:
    is_exportable = all(isinstance(item, exportable_types) for item in result)

if not is_exportable:
    fail(
        "result",
        error_type="InvalidResultType",
        error=f"'result' must be a CadQuery object or a non-empty list/tuple of CadQuery objects, got {type(result).__name__}.",
    )
    sys.exit(0)

try:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    cq.exporters.export(result, str(output_path))
except Exception as exc:
    fail("export", exc)
    sys.exit(0)

preview_error = None
try:
    preview_path.parent.mkdir(parents=True, exist_ok=True)
    cq.exporters.export(result, str(preview_path))
except Exception as exc:
    preview_error = str(exc)

payload = {
    "ok": True,
    "stage": "done",
    "output_path": str(output_path),
    "preview_path": str(preview_path) if preview_path.exists() else None,
    "warning": f"Failed to export STL preview: {preview_error}" if preview_error else None,
    "stdout": tail(stdout_buffer.getvalue()),
    "stderr": tail(stderr_buffer.getvalue()),
}
print(json.dumps(payload, ensure_ascii=False))
"""


def run_execute_cadquery(code: str, output_path: str = DEFAULT_CADQUERY_OUTPUT_PATH, prompt: str | None = None) -> str:
    try:
        if not code or not code.strip():
            return json_response({
                "ok": False,
                "stage": "input",
                "error_type": "EmptyCode",
                "error": "execute_cadquery requires non-empty CadQuery code.",
                "traceback_tail": "",
                "stdout": "",
                "stderr": "",
            })

        run_dir, output, run_id = resolve_run_output(output_path)
        preview = run_dir / "preview.stl"
        output_root = run_dir.parent.parent
    except Exception as exc:
        return json_response({
            "ok": False,
            "stage": "input",
            "error_type": type(exc).__name__,
            "error": str(exc),
            "traceback_tail": "",
            "stdout": "",
            "stderr": "",
        })

    request = {
        "code": code,
        "output_path": str(output),
        "preview_path": str(preview),
    }

    try:
        process = subprocess.run(
            [sys.executable, "-c", CADQUERY_RUNNER],
            input=json.dumps(request, ensure_ascii=False),
            cwd=WORKDIR,
            capture_output=True,
            text=True,
            timeout=120,
        )
    except subprocess.TimeoutExpired as exc:
        return json_response({
            "ok": False,
            "stage": "execution",
            "error_type": "TimeoutExpired",
            "error": "CadQuery execution timed out after 120 seconds.",
            "traceback_tail": "",
            "stdout": exc.stdout[-4000:] if exc.stdout else "",
            "stderr": exc.stderr[-4000:] if exc.stderr else "",
        })
    except Exception as exc:
        return json_response({
            "ok": False,
            "stage": "subprocess",
            "error_type": type(exc).__name__,
            "error": str(exc),
            "traceback_tail": "",
            "stdout": "",
            "stderr": "",
        })

    raw_output = process.stdout.strip()
    if not raw_output:
        return json_response({
            "ok": False,
            "stage": "subprocess",
            "error_type": "NoRunnerOutput",
            "error": "CadQuery runner returned no JSON output.",
            "traceback_tail": "",
            "stdout": "",
            "stderr": process.stderr[-4000:],
        })

    try:
        payload = json.loads(raw_output.splitlines()[-1])
    except Exception as exc:
        return json_response({
            "ok": False,
            "stage": "subprocess",
            "error_type": type(exc).__name__,
            "error": f"Failed to parse CadQuery runner output: {exc}",
            "traceback_tail": "",
            "stdout": raw_output[-4000:],
            "stderr": process.stderr[-4000:],
        })

    if process.returncode != 0 and payload.get("ok") is not False:
        payload = {
            "ok": False,
            "stage": "subprocess",
            "error_type": "RunnerFailed",
            "error": f"CadQuery runner exited with code {process.returncode}.",
            "traceback_tail": "",
            "stdout": raw_output[-4000:],
            "stderr": process.stderr[-4000:],
        }

    if payload.get("ok") and payload.get("output_path"):
        exported_path = Path(payload["output_path"])
        manifest_path = write_manifest(
            run_dir=run_dir,
            run_id=run_id,
            requested_output_path=output_path,
            output_path=exported_path,
            code=code,
            prompt=prompt,
            payload=payload,
        )
        latest_warning = update_latest_link(output_root, run_dir)

        payload["run_id"] = run_id
        payload["run_dir"] = workspace_relative(run_dir)
        payload["output_path"] = workspace_relative(exported_path)
        if payload.get("preview_path"):
            payload["preview_path"] = workspace_relative(Path(payload["preview_path"]))
        payload["manifest_path"] = workspace_relative(manifest_path)
        payload["latest_path"] = workspace_relative(output_root / "latest")
        if latest_warning:
            payload["warning"] = latest_warning

    return json_response(payload)


TOOL_SCHEMA = {
    "name": "execute_cadquery",
    "description": textwrap.dedent("""
        Execute CadQuery Python code and export the result as a STEP file.
        The code must assign the final CadQuery model to a variable named result.
        cadquery is pre-imported as cq.
        Do NOT supply output_path in normal usage. The tool automatically writes files to the configured artifact directory under a unique run directory. Only set output_path if the user explicitly requests a different location.
    """).strip(),
    "input_schema": {
        "type": "object",
        "properties": {
            "code": {
                "type": "string",
                "description": "CadQuery Python code. It must assign the final model to variable 'result'.",
            },
            "output_path": {
                "type": "string",
                "description": "Advanced: custom STEP file path or output directory. Omit this field to use the configured artifact directory (recommended).",
            },
        },
        "required": ["code"],
    },
}