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

"""Level / scene ops via UE EditorActorSubsystem — pure Python, no C++ needed.

Ops (dispatched by `entrypoint(payload_json)`):
  spawn_actor      {class, location:[x,y,z]?, rotation:[p,y,r]?, scale:[x,y,z]?, name?}
  destroy          {name}                       — by actor label or path
  get_selected     {}                            — list selected actors
  select           {names:[]}
  get_by_class     {class}
  find_by_name     {name}
  set_transform    {name, location?, rotation?, scale?}
  attach           {child, parent, socket?}
  save_level       {}
  load_level       {path}
  current_level    {}
"""
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 _load_class(p: str):
    if "/" in p:
        return unreal.load_class(None, p)
    for prefix in ("/Script/Engine.", "/Script/UnrealEd."):
        c = unreal.load_class(None, prefix + p)
        if c is not None:
            return c
    return None


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


def _rot(v, default=(0, 0, 0)):
    if v is None:
        v = default
    # rotator: pitch, yaw, roll
    return unreal.Rotator(float(v[0]), float(v[1]), float(v[2]))


def _find_actor(name: str):
    for a in _ess().get_all_level_actors():
        if not a:
            continue
        if a.get_actor_label() == name or a.get_name() == name or a.get_path_name() == name:
            return a
    return None


def spawn_actor(args: dict) -> dict:
    cls = _load_class(args["class"])
    if cls is None:
        return {"ok": False, "error": f"class not found: {args['class']!r}"}
    loc = _vec(args.get("location"))
    rot = _rot(args.get("rotation"))
    scale = _vec(args.get("scale"), (1, 1, 1))
    actor = _ess().spawn_actor_from_class(cls, loc, rot)
    if actor is None:
        return {"ok": False, "error": "spawn returned None"}
    actor.set_actor_scale3d(scale)
    if args.get("name"):
        actor.set_actor_label(args["name"])
    return {"ok": True, "name": actor.get_actor_label(), "path": actor.get_path_name()}


def destroy(args: dict) -> dict:
    a = _find_actor(args["name"])
    if not a:
        return {"ok": False, "error": "actor not found"}
    _ess().destroy_actor(a)
    return {"ok": True, "name": args["name"]}


def get_selected(args: dict) -> dict:
    return {"ok": True, "actors": [a.get_actor_label() for a in _ess().get_selected_level_actors()]}


def select(args: dict) -> dict:
    names = args.get("names") or []
    found = [a for a in (_find_actor(n) for n in names) if a]
    _ess().set_selected_level_actors(found)
    return {"ok": True, "selected": [a.get_actor_label() for a in found]}


def get_by_class(args: dict) -> dict:
    cls = _load_class(args["class"])
    if cls is None:
        return {"ok": False, "error": f"class not found: {args['class']!r}"}
    found = [a for a in _ess().get_all_level_actors() if isinstance(a, cls)]
    return {"ok": True, "actors": [a.get_actor_label() for a in found]}


def find_by_name(args: dict) -> dict:
    a = _find_actor(args["name"])
    if not a:
        return {"ok": False, "error": "not found"}
    return {"ok": True, "name": a.get_actor_label(), "path": a.get_path_name(), "class": a.get_class().get_name()}


def set_transform(args: dict) -> dict:
    a = _find_actor(args["name"])
    if not a:
        return {"ok": False, "error": "actor not found"}
    if args.get("location") is not None:
        a.set_actor_location(_vec(args["location"]), False, False)
    if args.get("rotation") is not None:
        a.set_actor_rotation(_rot(args["rotation"]), False)
    if args.get("scale") is not None:
        a.set_actor_scale3d(_vec(args["scale"], (1, 1, 1)))
    return {"ok": True}


def attach(args: dict) -> dict:
    child = _find_actor(args["child"])
    parent = _find_actor(args["parent"])
    if not child or not parent:
        return {"ok": False, "error": "child or parent not found"}
    socket = args.get("socket") or ""
    child.attach_to_actor(parent, socket, unreal.AttachmentRule.KEEP_WORLD,
                          unreal.AttachmentRule.KEEP_WORLD, unreal.AttachmentRule.KEEP_WORLD, False)
    return {"ok": True}


def save_level(args: dict) -> dict:
    les = unreal.get_editor_subsystem(unreal.LevelEditorSubsystem)
    ok = les.save_current_level()
    return {"ok": bool(ok)}


def load_level(args: dict) -> dict:
    les = unreal.get_editor_subsystem(unreal.LevelEditorSubsystem)
    ok = les.load_level(args["path"])
    return {"ok": bool(ok), "path": args["path"]}


def current_level(args: dict) -> dict:
    les = unreal.get_editor_subsystem(unreal.LevelEditorSubsystem)
    name = les.get_current_level_name() if hasattr(les, "get_current_level_name") else ""
    world = unreal.EditorLevelLibrary.get_editor_world() if hasattr(unreal, "EditorLevelLibrary") else None
    return {"ok": True, "name": name, "world": world.get_path_name() if world else None}


OPS = {
    "spawn_actor":   spawn_actor,
    "destroy":       destroy,
    "get_selected":  get_selected,
    "select":        select,
    "get_by_class":  get_by_class,
    "find_by_name":  find_by_name,
    "set_transform": set_transform,
    "attach":        attach,
    "save_level":    save_level,
    "load_level":    load_level,
    "current_level": current_level,
}


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}"}, {})
    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:  # 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-LEVEL] " + txt[:1500])
    except Exception:
        pass
    return txt