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

"""Read-only discovery ops via MCPGraphLibrary v0.3.

Ops:
  list_open_bps        {}                       — open BP editors
  list_graphs          {bp_path}                 — all graphs of a BP
  describe_bp          {bp_path}                 — parent/interfaces/vars/components/funcs/dispatchers
  get_selection        {bp_path}                 — selected nodes in BP editor
  resolve_call         {bp_path, function, guid} — CallFunction node → source bp/class
  find_var_refs        {bp_path, name}           — variable usage sites
  find_fn_refs         {bp_path, name}           — function call sites in this BP
  find_bp_by_parent    {parent_class_path}       — AssetRegistry search
  get_refs_of          {asset_path}              — assets referenced by this asset
  get_referencers      {asset_path}              — assets that reference this asset
  read_function        {bp_path, function}       — wrapper over apply_graph.read_graph for function
"""
from __future__ import annotations

import json
import os
import traceback

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


def _lib():
    L = getattr(unreal, "MCPGraphLibrary", None)
    if L is None:
        raise RuntimeError("MCPGraphLibrary missing — rebuild C++ plugin")
    return L


def _load_bp(path: str):
    a = unreal.EditorAssetLibrary.load_asset(path)
    if not isinstance(a, unreal.Blueprint):
        raise RuntimeError(f"not a blueprint: {path!r}")
    return a


def _find_graph(bp, name: str):
    L = _lib()
    if name in (None, "", "EventGraph"):
        return L.find_event_graph(bp)
    return L.find_function_graph(bp, name)


def list_open_bps(args: dict) -> dict:
    return {"ok": True, "open": list(_lib().list_open_blueprints())}


def list_graphs(args: dict) -> dict:
    bp = _load_bp(args["bp_path"])
    raw = _lib().list_graphs_json(bp)
    return {"ok": True, "bp_path": args["bp_path"], **json.loads(raw)}


def describe_bp(args: dict) -> dict:
    bp = _load_bp(args["bp_path"])
    raw = _lib().describe_blueprint_json(bp)
    return {"ok": True, **json.loads(raw)}


def get_selection(args: dict) -> dict:
    bp = _load_bp(args["bp_path"])
    raw = _lib().get_selected_nodes_json(bp)
    return {"ok": True, "bp_path": args["bp_path"], **json.loads(raw)}


def resolve_call(args: dict) -> dict:
    bp = _load_bp(args["bp_path"])
    g = _find_graph(bp, args.get("function") or "EventGraph")
    if g is None:
        return {"ok": False, "error": "graph not found"}
    raw = _lib().resolve_function_source_json(g, args["guid"])
    return {"ok": True, **json.loads(raw)}


def find_var_refs(args: dict) -> dict:
    bp = _load_bp(args["bp_path"])
    raw = _lib().find_variable_references_json(bp, args["name"])
    return {"ok": True, **json.loads(raw)}


def find_fn_refs(args: dict) -> dict:
    bp = _load_bp(args["bp_path"])
    raw = _lib().find_function_references_json(bp, args["name"])
    return {"ok": True, **json.loads(raw)}


def find_bp_by_parent(args: dict) -> dict:
    found = list(_lib().find_blueprints_by_parent(args["parent_class_path"]))
    return {"ok": True, "blueprints": found, "count": len(found)}


def get_refs_of(args: dict) -> dict:
    return {"ok": True, "references": list(_lib().get_referenced_assets(args["asset_path"]))}


def get_referencers(args: dict) -> dict:
    return {"ok": True, "referencers": list(_lib().get_referencers_of_asset(args["asset_path"]))}


def read_function(args: dict) -> dict:
    # Convenience: read_graph for a named function
    import importlib, apply_graph
    importlib.reload(apply_graph)
    return apply_graph.read_graph(args["bp_path"], args.get("function") or "EventGraph")


OPS = {
    "list_open_bps":     list_open_bps,
    "list_graphs":       list_graphs,
    "describe_bp":       describe_bp,
    "get_selection":     get_selection,
    "resolve_call":      resolve_call,
    "find_var_refs":     find_var_refs,
    "find_fn_refs":      find_fn_refs,
    "find_bp_by_parent": find_bp_by_parent,
    "get_refs_of":       get_refs_of,
    "get_referencers":   get_referencers,
    "read_function":     read_function,
}


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}"}, p={})
    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 + f"{rid}.json", "w", encoding="utf-8") as f:
                f.write(txt)
            unreal.log("[MCP-DISCOVER] " + txt[:1500])
    except Exception:
        pass
    return txt