unreal-engine-mcp-system-plugin/ue_side/capture.py

"""Headless capture ops — SEE editor content without opening any tab/window.

Backed by the C++ helper unreal.MCPCaptureLibrary (Live-Coding compile the
plugin if it's missing). Every op writes a PNG under Saved/<out_dir>/ and
returns its absolute path + file:// URL. The JS layer (capture_op) reads that
PNG back and attaches it INLINE so the agent actually sees the image.

Ops:
  asset      {asset_path, size?:[512,512], out_dir?:'MCP/capture'}
             Render ANY asset's thumbnail headlessly (mesh, material, texture,
             blueprint, niagara, anim, sound, data asset, ...). No editor open.
  thumbnail  alias for asset.
  material   {material_path, size?} — alias for asset, reads nicer in intent.
  mesh       {mesh_path, size?}     — alias for asset.
  scene      {location?:[x,y,z], rotation?:[pitch,yaw,roll], fov?:90,
              size?:[1280,720], out_dir?} — render the editor world from an
             arbitrary camera via an offscreen SceneCapture2D. No viewport.
  scene_from_actor {actor, distance?:300, pitch?:-20, yaw?:0, fov?:75, size?}
             — frame an existing level actor and capture it.
  list_assets {dir} — quick content-browser listing to find paths to capture.
"""
from __future__ import annotations

import json
import os
import traceback

try:
    import unreal  # type: ignore
except ImportError:
    unreal = None


def _cap_lib():
    L = getattr(unreal, "MCPCaptureLibrary", None)
    if L is None:
        raise RuntimeError(
            "MCPCaptureLibrary not exposed — Live-Coding compile the C++ plugin "
            "(live_coding_op compile), then retry."
        )
    return L


def _out_dir(args) -> tuple[str, str]:
    """Returns (abs_dir_with_os_sep, abs_dir_forward_slash)."""
    proj_saved = unreal.Paths.project_saved_dir()
    sub = (args.get("out_dir") or "MCP/capture").strip("/")
    fwd = proj_saved + sub + "/"
    os.makedirs(fwd, exist_ok=True)
    return fwd.rstrip("/").replace("/", os.sep), fwd


def _safe_name(path: str) -> str:
    return path.rsplit("/", 1)[-1].split(".")[0] or "capture"


def _ok_png(full_path: str, strategy: str, size) -> dict:
    norm = full_path.replace("\\", "/")
    return {
        "ok": True,
        "strategy": strategy,
        "png_path": full_path,
        "file_url": "file:///" + norm,
        "size": size,
    }


def _vec(v, d=(0.0, 0.0, 0.0)):
    if v is None:
        v = d
    return unreal.Vector(float(v[0]), float(v[1]), float(v[2]))


def _rot(v, d=(0.0, 0.0, 0.0)):
    if v is None:
        v = d
    # UE Rotator(pitch, yaw, roll); we accept [pitch, yaw, roll].
    return unreal.Rotator(float(v[0]), float(v[1]), float(v[2]))


def asset(args: dict) -> dict:
    path = args.get("asset_path") or args.get("material_path") or args.get("mesh_path") or args.get("path")
    if not path:
        return {"ok": False, "error": "asset_path required"}
    if not unreal.EditorAssetLibrary.does_asset_exist(path):
        return {"ok": False, "error": f"asset does not exist: {path}"}
    size = args.get("size") or [512, 512]
    w, h = int(size[0]), int(size[1])
    os_dir, _ = _out_dir(args)
    name = _safe_name(path)
    _ret = _cap_lib().render_asset_thumbnail_to_png(path, os_dir, name, w, h)
    if isinstance(_ret, (tuple, list)):
        if len(_ret) >= 3:
            ok, full, err = _ret[0], _ret[1], _ret[2]
        elif len(_ret) == 2:
            full, err = _ret[0], _ret[1]
            ok = bool(full)
        else:
            ok, full, err = bool(_ret[0]), (_ret[0] if isinstance(_ret[0], str) else ""), ""
    else:
        full = _ret
        ok = bool(full)
        err = ""
    if not ok:
        return {"ok": False, "error": err or "thumbnail render failed", "asset_path": path}
    res = _ok_png(full, "AssetThumbnail", [w, h])
    res["asset_path"] = path
    return res


def scene(args: dict) -> dict:
    size = args.get("size") or [1280, 720]
    w, h = int(size[0]), int(size[1])
    os_dir, _ = _out_dir(args)
    name = args.get("name") or "scene"
    loc = _vec(args.get("location"))
    rot = _rot(args.get("rotation"))
    fov = float(args.get("fov", 90.0))
    ok, full, err = _cap_lib().capture_editor_scene_to_png(loc, rot, fov, os_dir, name, w, h)
    if not ok:
        return {"ok": False, "error": err or "scene capture failed"}
    res = _ok_png(full, "SceneCapture2D", [w, h])
    res["camera"] = {"location": [loc.x, loc.y, loc.z],
                     "rotation": [rot.pitch, rot.yaw, rot.roll], "fov": fov}
    return res


def scene_from_actor(args: dict) -> dict:
    name = args.get("actor")
    if not name:
        return {"ok": False, "error": "actor required"}
    ess = unreal.get_editor_subsystem(unreal.EditorActorSubsystem)
    target = None
    for a in ess.get_all_level_actors():
        if a and (a.get_actor_label() == name or a.get_name() == name):
            target = a
            break
    if target is None:
        return {"ok": False, "error": f"actor not found: {name}"}

    origin, extent = target.get_actor_bounds(False)
    radius = max(extent.x, extent.y, extent.z, 50.0)
    pitch = float(args.get("pitch", -20.0))
    yaw = float(args.get("yaw", 0.0))
    fov = float(args.get("fov", 75.0))
    dist = float(args.get("distance", radius * 3.0))

    # Place camera 'dist' away along yaw/pitch, looking back at the actor center.
    look_dir = unreal.MathLibrary.get_forward_vector(unreal.Rotator(pitch, yaw, 0.0))
    cam_loc = origin - (look_dir * dist)
    look_rot = unreal.MathLibrary.find_look_at_rotation(cam_loc, origin)

    size = args.get("size") or [1280, 720]
    w, h = int(size[0]), int(size[1])
    os_dir, _ = _out_dir(args)
    fname = args.get("name") or f"actor_{_safe_name(name)}"
    ok, full, err = _cap_lib().capture_editor_scene_to_png(cam_loc, look_rot, fov, os_dir, fname, w, h)
    if not ok:
        return {"ok": False, "error": err or "scene capture failed"}
    res = _ok_png(full, "SceneCapture2D/actor", [w, h])
    res["actor"] = name
    res["camera"] = {"location": [cam_loc.x, cam_loc.y, cam_loc.z], "fov": fov}
    return res


def list_assets(args: dict) -> dict:
    folder = args.get("dir") or "/Game"
    recursive = bool(args.get("recursive", False))
    paths = unreal.EditorAssetLibrary.list_assets(folder, recursive=recursive, include_folder=False)
    return {"ok": True, "dir": folder, "count": len(paths), "assets": list(paths)[:500]}


OPS = {
    "asset": asset,
    "thumbnail": asset,
    "material": asset,
    "mesh": asset,
    "scene": scene,
    "scene_from_actor": scene_from_actor,
    "list_assets": list_assets,
}


def entrypoint(payload_json: str) -> str:
    try:
        p = json.loads(payload_json)
    except Exception as e:  # noqa: BLE001
        return _tee({"ok": False, "error": f"bad json: {e}"}, {})
    fn = OPS.get(p.get("op"))
    if not fn:
        return _tee({"ok": False, "error": f"unknown op {p.get('op')!r}", "available": list(OPS)}, p)
    try:
        return _tee(fn(p), p)
    except Exception as e:  # noqa: BLE001
        return _tee({"ok": False, "error": str(e), "traceback": traceback.format_exc()}, p)


def _tee(result: dict, p: dict) -> str:
    txt = json.dumps(result)
    try:
        if unreal is not None:
            sd = unreal.Paths.project_saved_dir() + "MCP/"
            os.makedirs(sd, exist_ok=True)
            rid = p.get("__rid") or "last"
            with open(sd + "last.json", "w", encoding="utf-8") as f:
                f.write(txt)
            if rid != "last":
                with open(sd + f"{rid}.json", "w", encoding="utf-8") as f:
                    f.write(txt)
            unreal.log("[MCP-CAPTURE] " + txt[:1500])
    except Exception:
        pass
    return txt