mcp-ue/ue_mcp_server.py

"""MCP server exposing UE documentation tools at item granularity."""

import os
import re
import subprocess
from pathlib import Path

from mcp.server.fastmcp import FastMCP

mcp = FastMCP("ue-docs")


# ---------------------------------------------------------------------------
# Env helpers
# ---------------------------------------------------------------------------

def _docs_root() -> Path:
    p = os.environ.get("UE_DOCS_PATH", "")
    if not p:
        raise RuntimeError("UE_DOCS_PATH environment variable not set")
    return Path(p)


def _engine_root() -> Path:
    p = os.environ.get("UE_ENGINE_ROOT", "")
    if not p:
        raise RuntimeError("UE_ENGINE_ROOT environment variable not set")
    return Path(p)


_type_index_cache: dict[str, str] = {}
_type_index_mtime: float = 0.0


def _load_type_index() -> dict[str, str]:
    """Returns {TypeName: relative_md_path}, reloading only if the file has changed."""
    global _type_index_cache, _type_index_mtime
    idx_path = _docs_root() / "type-index.txt"
    mtime = idx_path.stat().st_mtime
    if mtime != _type_index_mtime:
        result = {}
        for line in idx_path.read_text().splitlines():
            if ": " in line:
                name, path = line.split(": ", 1)
                result[name.strip()] = path.strip()
        _type_index_cache = result
        _type_index_mtime = mtime
    return _type_index_cache


# ---------------------------------------------------------------------------
# Section extraction helpers
# ---------------------------------------------------------------------------

def _extract_class_section(text: str, class_name: str) -> str:
    """Return the markdown section for class_name (### `ClassName` … next ### `)."""
    # Match the heading for this class
    pattern = rf'^### `{re.escape(class_name)}`'
    m = re.search(pattern, text, re.MULTILINE)
    if not m:
        return ""
    start = m.start()
    # Find the next same-level heading
    next_h3 = re.search(r'^### `', text[m.end():], re.MULTILINE)
    if next_h3:
        end = m.end() + next_h3.start()
    else:
        end = len(text)
    return text[start:end]


def _build_overview(section: str) -> str:
    """Compact overview: header block + name lists for Properties and Functions."""
    lines = section.splitlines()

    prop_names: list[str] = []
    func_names: list[str] = []  # deduped
    seen_funcs: set[str] = set()

    header_lines: list[str] = []
    in_props = False
    in_funcs = False

    for line in lines:
        if line.startswith("#### Properties"):
            in_props = True
            in_funcs = False
            continue
        if line.startswith("#### Functions"):
            in_funcs = True
            in_props = False
            continue
        if line.startswith("#### ") or line.startswith("### "):
            # Some other subsection — stop collecting
            in_props = False
            in_funcs = False

        if in_props:
            m = re.match(r"- `(\w+)`", line)
            if m:
                prop_names.append(m.group(1))
        elif in_funcs:
            m = re.match(r"##### `(\w+)\(", line)
            if m:
                name = m.group(1)
                if name not in seen_funcs:
                    seen_funcs.add(name)
                    func_names.append(name)
        else:
            header_lines.append(line)

    # Trim trailing blank lines from header
    while header_lines and not header_lines[-1].strip():
        header_lines.pop()

    parts = ["\n".join(header_lines)]
    if prop_names:
        parts.append(f"\nProperties ({len(prop_names)}): {', '.join(prop_names)}")
    if func_names:
        parts.append(f"Functions ({len(func_names)}): {', '.join(func_names)}")

    return "\n".join(parts)


def _extract_member(section: str, member_name: str) -> str:
    """Return the full doc entry for member_name (all overloads for functions)."""
    escaped = re.escape(member_name)

    # Try function first (##### `Name(`)
    pattern = rf'^##### `{escaped}\('
    m = re.search(pattern, section, re.MULTILINE)
    if m:
        start = m.start()
        # Collect until the next ##### ` heading
        next_h5 = re.search(r'^##### `', section[m.end():], re.MULTILINE)
        if next_h5:
            end = m.end() + next_h5.start()
        else:
            end = len(section)
        return section[start:end].rstrip()

    # Try property (- `Name`)
    m = re.search(rf'^- `{escaped}`', section, re.MULTILINE)
    if m:
        return section[m.start():section.index("\n", m.start())].rstrip()

    return ""


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------

@mcp.tool()
def search_types(pattern: str) -> str:
    """Search type-index.txt for UE types whose name matches pattern (case-insensitive).
    Returns matching 'TypeName: path/to/File.md' lines (up to 50).
    Use this to locate a type before calling get_class_overview or get_file."""
    idx = _docs_root() / "type-index.txt"
    lines = [l for l in idx.read_text().splitlines()
             if re.search(pattern, l, re.IGNORECASE)]
    return "\n".join(lines[:50]) or "(no matches)"


@mcp.tool()
def get_class_overview(class_name: str) -> str:
    """Compact overview of a UE class/struct: description, base classes,
    and flat lists of property and function names. Use get_member() for
    full signatures and docs. Use get_file() if you need delegates/enums too."""
    index = _load_type_index()
    if class_name not in index:
        return f"'{class_name}' not in type index. Try search_types('{class_name}')."
    text = (_docs_root() / index[class_name]).read_text()
    section = _extract_class_section(text, class_name)
    if not section:
        return f"Section for '{class_name}' not found in {index[class_name]}."
    return _build_overview(section)


@mcp.tool()
def get_member(class_name: str, member_name: str) -> str:
    """Full documentation entry for one function or property in a UE class.
    For overloaded functions, returns all overloads together."""
    index = _load_type_index()
    if class_name not in index:
        return f"'{class_name}' not found."
    text = (_docs_root() / index[class_name]).read_text()
    section = _extract_class_section(text, class_name)
    if not section:
        return f"'{class_name}' section missing."
    entry = _extract_member(section, member_name)
    return entry or f"'{member_name}' not found in {class_name}."


@mcp.tool()
def get_file(relative_path: str) -> str:
    """Full content of a documentation .md file.
    relative_path is relative to $UE_DOCS_PATH, as returned by search_types().
    Use this when you need delegates, top-level enums, or multiple classes from one file."""
    path = _docs_root() / relative_path
    if not path.exists():
        return f"File not found: {relative_path}"
    return path.read_text()


@mcp.tool()
def search_source(pattern: str, path_hint: str = "") -> str:
    """Grep UE source .h files for pattern. Requires $UE_ENGINE_ROOT.
    path_hint: optional subdirectory under Engine/Source/ (e.g. 'Runtime/Engine').
    Searches both Engine/Source/ and Engine/Plugins/.
    Returns up to 40 matching lines with file:line context."""
    engine = _engine_root()
    candidates = [engine / "Source", engine / "Plugins"]
    search_dirs = [d / path_hint if path_hint else d for d in candidates]
    search_dirs = [d for d in search_dirs if d.exists()]
    if not search_dirs:
        return "(no matches)"
    grep = subprocess.Popen(
        ["grep", "-rn", "--include=*.h", pattern, *map(str, search_dirs)],
        stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True
    )
    lines = []
    for line in grep.stdout:
        lines.append(line)
        if len(lines) == 40:
            grep.kill()
            break
    grep.wait()
    return "".join(lines) or "(no matches)"


if __name__ == "__main__":
    mcp.run()