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

"""Render ops — PostProcess volumes, Render targets, Scene captures, viewport screenshots.

Ops:
  spawn_post_process_volume {location?, unbound?:true, settings?:{key:value}}
  set_pp_setting            {actor_name, key, value}    — sets PostProcessVolume.Settings entries
  list_pp_volumes           {}
  create_render_target_2d   {path, width?:512, height?:512, format?:'RGBA16f|RGBA8'}
  spawn_scene_capture_2d    {target_render_target_path, location?, rotation?, name?}
  take_screenshot           {filename?, resolution?, hdr?:false}      — uses HighResShot console cmd
  set_viewport_realtime     {enabled:bool}
  capture_thumbnail         {asset_path}                — refresh editor thumbnail
"""
from __future__ import annotations

import json
import os
import traceback

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


def _ess(): return unreal.get_editor_subsystem(unreal.EditorActorSubsystem)
def _world(): return unreal.EditorLevelLibrary.get_editor_world()
def _vec(v, d=(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)):
    if v is None: v = d
    return unreal.Rotator(float(v[0]), float(v[1]), float(v[2]))
def _find(name):
    for a in _ess().get_all_level_actors():
        if a and (a.get_actor_label() == name or a.get_name() == name):
            return a
    return None


def spawn_post_process_volume(args):
    cls = unreal.load_class(None, "/Script/Engine.PostProcessVolume")
    actor = _ess().spawn_actor_from_class(cls, _vec(args.get("location")), _rot(None))
    if actor is None:
        return {"ok": False, "error": "spawn failed"}
    if args.get("unbound", True):
        actor.set_editor_property("unbound", True)
    settings = args.get("settings") or {}
    if settings:
        pp = actor.get_editor_property("settings")
        for k, v in settings.items():
            try:
                pp.set_editor_property(k, v)
            except Exception:
                pass
        actor.set_editor_property("settings", pp)
    label = args.get("name") or "PostProcessVolume_MCP"
    actor.set_actor_label(label)
    return {"ok": True, "name": label, "path": actor.get_path_name()}


def set_pp_setting(args):
    a = _find(args["actor_name"])
    if a is None:
        return {"ok": False, "error": "actor not found"}
    try:
        pp = a.get_editor_property("settings")
        pp.set_editor_property(args["key"], args["value"])
        a.set_editor_property("settings", pp)
        return {"ok": True}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def list_pp_volumes(args):
    cls = unreal.load_class(None, "/Script/Engine.PostProcessVolume")
    out = []
    for a in _ess().get_all_level_actors():
        if a and a.get_class() == cls.get_default_object().get_class():
            out.append({"name": a.get_actor_label(), "unbound": bool(a.get_editor_property("unbound"))})
    return {"ok": True, "volumes": out}


def create_render_target_2d(args):
    folder, name = args["path"].rsplit("/", 1)
    factory = unreal.CanvasRenderTarget2DFactoryNew() if hasattr(unreal, "CanvasRenderTarget2DFactoryNew") else None
    asset_cls = unreal.TextureRenderTarget2D
    factory = unreal.TextureRenderTargetFactoryNew() if hasattr(unreal, "TextureRenderTargetFactoryNew") else factory
    if factory is None:
        return {"ok": False, "error": "no RT factory exposed"}
    at = unreal.AssetToolsHelpers.get_asset_tools()
    asset = at.create_asset(name, folder, asset_cls, factory)
    if asset is None:
        return {"ok": False, "error": "create returned None"}
    asset.set_editor_property("size_x", int(args.get("width") or 512))
    asset.set_editor_property("size_y", int(args.get("height") or 512))
    fmt = args.get("format") or "RGBA8"
    fmt_enum = unreal.TextureRenderTargetFormat.RTF_RGBA16F if fmt == "RGBA16f" else unreal.TextureRenderTargetFormat.RTF_RGBA8
    try:
        asset.set_editor_property("render_target_format", fmt_enum)
    except Exception:
        pass
    unreal.EditorAssetLibrary.save_asset(asset.get_path_name())
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def spawn_scene_capture_2d(args):
    rt = unreal.EditorAssetLibrary.load_asset(args["target_render_target_path"])
    if rt is None:
        return {"ok": False, "error": "render target not found"}
    cls = unreal.load_class(None, "/Script/Engine.SceneCapture2D")
    actor = _ess().spawn_actor_from_class(cls, _vec(args.get("location")), _rot(args.get("rotation")))
    if actor is None:
        return {"ok": False, "error": "spawn failed"}
    try:
        comp = actor.get_component_by_class(unreal.SceneCaptureComponent2D)
        comp.set_editor_property("texture_target", rt)
    except Exception as e:
        return {"ok": False, "error": f"set RT failed: {e}"}
    label = args.get("name") or "SceneCapture2D_MCP"
    actor.set_actor_label(label)
    return {"ok": True, "name": label}


def take_screenshot(args):
    res = args.get("resolution") or "1920x1080"
    fname = args.get("filename") or ""
    hdr = "1" if args.get("hdr") else ""
    cmd = f"HighResShot {res} {fname} {hdr}".strip()
    try:
        unreal.SystemLibrary.execute_console_command(_world(), cmd)
        return {"ok": True, "command": cmd, "note": "Saved to <Project>/Saved/Screenshots/<Platform>/"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def set_viewport_realtime(args):
    try:
        ues = unreal.get_editor_subsystem(unreal.UnrealEditorSubsystem)
        # No direct API; use console command
        cmd = "r.Realtime 1" if args.get("enabled") else "r.Realtime 0"
        unreal.SystemLibrary.execute_console_command(_world(), cmd)
        return {"ok": True, "command": cmd}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def capture_thumbnail(args):
    paths = unreal.Array(str)
    paths.append(args["asset_path"])
    try:
        unreal.AssetEditorSubsystem if False else None
        # No direct python wrapper; rely on AssetTools? Fall back to editor cmd
        unreal.SystemLibrary.execute_console_command(_world(), f"AssetThumbnailCapture {args['asset_path']}")
        return {"ok": True, "note": "issued console command"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


OPS = {
    "spawn_post_process_volume": spawn_post_process_volume,
    "set_pp_setting":            set_pp_setting,
    "list_pp_volumes":           list_pp_volumes,
    "create_render_target_2d":   create_render_target_2d,
    "spawn_scene_capture_2d":    spawn_scene_capture_2d,
    "take_screenshot":           take_screenshot,
    "set_viewport_realtime":     set_viewport_realtime,
    "capture_thumbnail":         capture_thumbnail,
}


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


def _tee(result, p):
    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-RENDER] " + txt[:1500])
    except Exception:
        pass
    return txt