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

"""Asset CRUD via UE AssetTools — pure Python, no C++ needed.

Exposed ops (dispatched by `entrypoint(payload_json)` from the MCP server):
  create_blueprint     {path, parent_class}
  create_widget_bp     {path}
  create_enum          {path, values:[]}
  create_struct        {path, fields:[{name,type}]}
  create_data_asset    {path, class}
  create_data_table    {path, row_struct}
  duplicate            {src, dst}
  rename               {src, dst}
  move                 {src, dst_folder}
  delete               {path}
  reparent_blueprint   {bp, new_parent_class}
"""
from __future__ import annotations

import json
import os
import traceback

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


def _at():
    return unreal.AssetToolsHelpers.get_asset_tools()


def _load_class(path_or_name: str):
    if "/" in path_or_name:
        return unreal.load_class(None, path_or_name)
    for prefix in ("/Script/Engine.", "/Script/CoreUObject.", "/Script/UMG.",
                   "/Script/UnrealEd.", "/Script/Kismet."):
        c = unreal.load_class(None, prefix + path_or_name)
        if c is not None:
            return c
    return None


def _split_pkg(path: str):
    """/Game/Foo/Bar/Asset -> (/Game/Foo/Bar, Asset)"""
    if "/" not in path:
        raise RuntimeError(f"bad asset path {path!r}")
    folder, name = path.rsplit("/", 1)
    return folder, name


def create_blueprint(args: dict) -> dict:
    path = args["path"]
    parent = _load_class(args.get("parent_class") or "Actor")
    if parent is None:
        return {"ok": False, "error": f"parent class not found: {args.get('parent_class')!r}"}
    folder, name = _split_pkg(path)
    factory = unreal.BlueprintFactory()
    factory.set_editor_property("parent_class", parent)
    asset = _at().create_asset(name, folder, unreal.Blueprint, factory)
    if asset is None:
        return {"ok": False, "error": "create_asset returned None"}
    unreal.EditorAssetLibrary.save_loaded_asset(asset)
    return {"ok": True, "path": asset.get_path_name().split(".")[0], "parent": parent.get_name()}


def create_widget_bp(args: dict) -> dict:
    path = args["path"]
    folder, name = _split_pkg(path)
    factory = unreal.WidgetBlueprintFactory()
    asset = _at().create_asset(name, folder, unreal.WidgetBlueprint, factory)
    if asset is None:
        return {"ok": False, "error": "create_asset returned None"}
    unreal.EditorAssetLibrary.save_loaded_asset(asset)
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def create_enum(args: dict) -> dict:
    path = args["path"]
    values = args.get("values") or []
    folder, name = _split_pkg(path)
    factory = unreal.EnumFactory()
    asset = _at().create_asset(name, folder, unreal.UserDefinedEnum, factory)
    if asset is None:
        return {"ok": False, "error": "create_asset returned None"}
    for i, v in enumerate(values):
        try:
            unreal.EnumEditorUtils.add_enumerator_to_user_defined_enum(asset, str(v))
        except Exception:
            pass
    unreal.EditorAssetLibrary.save_loaded_asset(asset)
    return {"ok": True, "path": asset.get_path_name().split(".")[0], "values_added": len(values)}


def create_struct(args: dict) -> dict:
    path = args["path"]
    fields = args.get("fields") or []  # [{name, type}]
    folder, name = _split_pkg(path)
    factory = unreal.StructureFactory()
    asset = _at().create_asset(name, folder, unreal.UserDefinedStruct, factory)
    if asset is None:
        return {"ok": False, "error": "create_asset returned None"}
    # Field editing via StructureEditorUtils not exposed reliably in stock Python.
    # User can rename default member via UI; we just create the struct.
    unreal.EditorAssetLibrary.save_loaded_asset(asset)
    return {"ok": True, "path": asset.get_path_name().split(".")[0], "note": "Field add not exposed in stock Python; edit in editor or pre-create C++ struct."}


def create_data_asset(args: dict) -> dict:
    path = args["path"]
    cls = _load_class(args.get("class") or "")
    if cls is None:
        return {"ok": False, "error": f"class not found: {args.get('class')!r}"}
    folder, name = _split_pkg(path)
    factory = unreal.DataAssetFactory()
    factory.set_editor_property("data_asset_class", cls)
    asset = _at().create_asset(name, folder, cls, factory)
    if asset is None:
        return {"ok": False, "error": "create_asset returned None"}
    unreal.EditorAssetLibrary.save_loaded_asset(asset)
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def create_data_table(args: dict) -> dict:
    path = args["path"]
    row_struct_path = args.get("row_struct") or ""
    struct = unreal.load_object(None, row_struct_path) if row_struct_path else None
    if struct is None:
        return {"ok": False, "error": f"row_struct not found: {row_struct_path!r}"}
    folder, name = _split_pkg(path)
    factory = unreal.DataTableFactory()
    factory.set_editor_property("struct", struct)
    asset = _at().create_asset(name, folder, unreal.DataTable, factory)
    if asset is None:
        return {"ok": False, "error": "create_asset returned None"}
    unreal.EditorAssetLibrary.save_loaded_asset(asset)
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def duplicate(args: dict) -> dict:
    src, dst = args["src"], args["dst"]
    asset = unreal.EditorAssetLibrary.duplicate_asset(src, dst)
    return {"ok": asset is not None, "src": src, "dst": dst}


def rename(args: dict) -> dict:
    src, dst = args["src"], args["dst"]
    ok = unreal.EditorAssetLibrary.rename_asset(src, dst)
    return {"ok": bool(ok), "src": src, "dst": dst}


def move(args: dict) -> dict:
    src, folder = args["src"], args["dst_folder"]
    name = src.rsplit("/", 1)[-1]
    dst = folder.rstrip("/") + "/" + name
    ok = unreal.EditorAssetLibrary.rename_asset(src, dst)
    return {"ok": bool(ok), "src": src, "dst": dst}


def delete(args: dict) -> dict:
    path = args["path"]
    ok = unreal.EditorAssetLibrary.delete_asset(path)
    return {"ok": bool(ok), "path": path}


def reparent_blueprint(args: dict) -> dict:
    bp_path = args["bp"]
    new_parent = _load_class(args["new_parent_class"])
    if new_parent is None:
        return {"ok": False, "error": f"new_parent_class not found: {args['new_parent_class']!r}"}
    bp = unreal.EditorAssetLibrary.load_asset(bp_path)
    if not isinstance(bp, unreal.Blueprint):
        return {"ok": False, "error": "not a blueprint"}
    # WidgetBlueprint (and others) don't expose ParentClass as a settable editor
    # property — use the proper reparent API, fall back to set_editor_property.
    try:
        unreal.BlueprintEditorLibrary.reparent_blueprint(bp, new_parent)
    except Exception:
        bp.set_editor_property("parent_class", new_parent)
    unreal.BlueprintEditorLibrary.compile_blueprint(bp)
    unreal.EditorAssetLibrary.save_loaded_asset(bp)
    return {"ok": True, "bp": bp_path, "new_parent": new_parent.get_name()}


OPS = {
    "create_blueprint":   create_blueprint,
    "create_widget_bp":   create_widget_bp,
    "create_enum":        create_enum,
    "create_struct":      create_struct,
    "create_data_asset":  create_data_asset,
    "create_data_table":  create_data_table,
    "duplicate":          duplicate,
    "rename":             rename,
    "move":               move,
    "delete":             delete,
    "reparent_blueprint": reparent_blueprint,
}

# Member-variable authoring lives in variables.py and is surfaced as the
# dedicated `variable_op` tool. Merge its NON-colliding op names here too so the
# capability is reachable via `asset_op` without a JS server restart. (We skip
# the short aliases add/rename/delete/list — they would shadow asset ops.)
try:
    import importlib as _il
    import variables as _vars
    _il.reload(_vars)
    for _k in ("add_variable", "set_var_default", "set_var_meta",
               "rename_variable", "delete_variable", "list_variables"):
        OPS[_k] = _vars.VAR_OPS[_k]
except Exception:
    pass


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-ASSETS] " + txt[:1500])
    except Exception:
        pass
    return txt