ayon-mcp/src/ayon_mcp/server.py

"""AYON MCP Server — main entry point.

Usage:
    AYON_URL=http://172.18.0.2:5000 \\
    AYON_KEY=sk-... \\
    ayon-mcp

Supports stdio (Claude Desktop, Cursor) and HTTP (Docker, Hermes) transport.
Set AYON_TRANSPORT=http for HTTP mode.
"""

import json
import os
import sys
from typing import Optional

from mcp.server.fastmcp import FastMCP
from mcp.server.transport_security import TransportSecuritySettings

from .client import AyonClient

# ── Bootstrap ────────────────────────────────────────────

app = FastMCP(
    "ayon-publish",
    transport_security=TransportSecuritySettings(
        enable_dns_rebinding_protection=False
    ),
)
client = AyonClient()


def _ok(**kwargs) -> str:
    return json.dumps(dict(status="ok", **kwargs), indent=2, ensure_ascii=False)


def _err(msg: str) -> str:
    return json.dumps({"error": msg}, indent=2)


# ── Discovery Tools ──────────────────────────────────────

@app.tool()
async def ayon_list_projects() -> str:
    """List all available AYON projects."""
    try:
        projects = client.list_projects()
        summary = [{"name": p["name"], "code": p.get("code")} for p in projects]
        return json.dumps(summary, indent=2)
    except Exception as e:
        return _err(str(e))


@app.tool()
async def ayon_list_folders(project: str) -> str:
    """List folders in an AYON project.

    Args:
        project: Project name (e.g. 'MyProject')
    """
    try:
        folders = client.list_folders(project)
        summary = [
            {"path": f.get("path"), "type": f.get("folderType"), "id": f.get("id")}
            for f in folders
        ]
        return json.dumps(summary, indent=2)
    except Exception as e:
        return _err(str(e))


@app.tool()
async def ayon_list_products(project: str, folder_path: str) -> str:
    """List published products in a folder.

    Uses GraphQL for AYON >=1.15 compatibility.
    Accepts paths with or without leading slash.

    Args:
        project: Project name
        folder_path: Folder path (e.g. '/assets/chair' or 'assets/chair')
    """
    try:
        folder = client.find_folder(project, folder_path)
        if not folder:
            return _err(f"Folder '{folder_path}' not found in '{project}'")
        products = client.list_products(project, folder["id"])
        summary = [
            {"name": p["name"], "type": p.get("productType"), "id": p["id"]}
            for p in products
        ]
        return json.dumps(summary, indent=2)
    except Exception as e:
        return _err(str(e))


@app.tool()
async def ayon_get_product(project: str, folder_path: str, product_name: str) -> str:
    """Get details about a specific product, including latest version.

    Uses GraphQL for AYON >=1.15 compatibility.

    Args:
        project: Project name
        folder_path: Folder path
        product_name: Product name to find
    """
    try:
        folder = client.find_folder(project, folder_path)
        if not folder:
            return _err(f"Folder '{folder_path}' not found")
        product = client.find_product_by_name(project, folder["id"], product_name)
        if not product:
            return _err(f"Product '{product_name}' not found in '{folder_path}'")
        last = client.get_last_version(project, product["id"])
        return json.dumps({
            "name": product["name"],
            "type": product.get("productType"),
            "id": product["id"],
            "latest_version": last["version"] if last else None,
        }, indent=2)
    except Exception as e:
        return _err(str(e))


# ── Folder Management ────────────────────────────────────

@app.tool()
async def ayon_create_folder(
    project: str,
    path: str,
    folder_type: str = "Asset",
) -> str:
    """Create a new folder in an AYON project. Creates parent folders as needed.

    Args:
        project: Project name (e.g. 'lumenfjord')
        path: Folder path, slash-separated (e.g. 'assets/test/my_folder')
        folder_type: AYON folder type: 'Asset', 'Folder', 'Shot', 'Sequence',
                     'Episode', 'Library' (default: 'Asset')
    """
    VALID_TYPES = {"Folder", "Library", "Asset", "Episode", "Sequence", "Shot"}
    if folder_type not in VALID_TYPES:
        return _err(f"Invalid folder type '{folder_type}'. Must be one of: {', '.join(sorted(VALID_TYPES))}")

    try:
        folder = client.find_or_create_folder(project, path, folder_type)
        return json.dumps({
            "name": folder["name"],
            "path": folder.get("path"),
            "id": folder["id"],
            "created": True,
        }, indent=2)
    except Exception as e:
        return _err(str(e))


# ── Publishing Tools ─────────────────────────────────────

def _publish_product(
    project: str,
    folder_path: str,
    product_name: str,
    product_type: str,
    representation_name: str,
    files_list: list[dict],
    variant: str = "",
) -> str:
    """Core publish logic — creates product or auto-versions existing one.

    Returns a JSON string with status and details.
    The caller decides if this is modelMain or imageMain.
    """
    folder = client.find_folder(project, folder_path)
    if not folder:
        return _err(f"Folder '{folder_path}' not found in '{project}'")

    full_name = product_name
    if variant:
        full_name = f"{product_name}_{variant}"

    # Check if product already exists
    existing = client.find_product_by_name(project, folder["id"], full_name)

    if existing:
        # Product exists — auto-version
        last = client.get_last_version(project, existing["id"])
        version_num = (last["version"] + 1) if last else 1

        version = client.create_version(
            project=project,
            product_id=existing["id"],
            version=version_num,
        )

        rep = client.create_representation(
            project=project,
            version_id=version["id"],
            name=representation_name,
            files=files_list,
        )

        return _ok(
            action="versioned_existing",
            product=full_name,
            type=product_type,
            version=version_num,
            existing_product=True,
            product_id=existing["id"],
            version_id=version["id"],
            folder=folder_path,
            files=[f["name"] for f in files_list],
            note=f"Product '{full_name}' already existed. Created version {version_num}.",
        )
    else:
        # New product
        product = client.create_product(
            project=project,
            name=full_name,
            product_type=product_type,
            folder_id=folder["id"],
        )

        last = client.get_last_version(project, product["id"])
        version_num = (last["version"] + 1) if last else 1

        version = client.create_version(
            project=project,
            product_id=product["id"],
            version=version_num,
        )

        rep = client.create_representation(
            project=project,
            version_id=version["id"],
            name=representation_name,
            files=files_list,
        )

        return _ok(
            action="created_new",
            product=full_name,
            type=product_type,
            version=version_num,
            existing_product=False,
            product_id=product["id"],
            version_id=version["id"],
            folder=folder_path,
            files=[f["name"] for f in files_list],
            note=f"New product '{full_name}' created with version {version_num}.",
        )


@app.tool()
async def ayon_publish_texture_set(
    project: str,
    folder_path: str,
    product_name: str,
    variant: str = "",
    textures: str = "[]",
) -> str:
    """Publish a texture set (multiple maps) as a single product.

    If a product with the same name already exists, a new VERSION is created
    instead of a duplicate product. The response indicates which happened.

    Args:
        project: Project name (e.g. 'DemoProject')
        folder_path: Folder path (e.g. '/assets/chair' or 'assets/chair')
        product_name: Base product name (e.g. 'texture_chair')
        variant: Optional variant (e.g. 'brown_leather')
        textures: JSON array of texture info:
            [{"name": "BaseColor.png", "size": 1048576}, ...]
    """
    try:
        textures_list = json.loads(textures)
        return _publish_product(
            project=project,
            folder_path=folder_path,
            product_name=product_name,
            product_type="imageMain",
            representation_name="texture_set",
            files_list=textures_list,
            variant=variant,
        )
    except Exception as e:
        return _err(str(e))


@app.tool()
async def ayon_publish_model(
    project: str,
    folder_path: str,
    product_name: str,
    model_type: str = "abc",
    variant: str = "",
    files: str = "[]",
) -> str:
    """Publish a 3D model (ABC, USD, FBX, etc.).

    If a product with the same name already exists, a new VERSION is created
    instead of a duplicate product. The response indicates which happened
    via the 'action' field ('created_new' or 'versioned_existing').

    Args:
        project: Project name
        folder_path: Folder path (e.g. '/assets/chair' or 'assets/chair')
        product_name: Product name (e.g. 'modelMain')
        model_type: File format: 'abc', 'usd', 'fbx', 'obj', 'blend'
        variant: Optional variant (e.g. 'highpoly')
        files: JSON array of file info:
            [{"name": "chair.abc", "size": 2048000}]
    """
    VALID_FORMATS = {"abc", "usd", "fbx", "obj", "blend"}
    if model_type not in VALID_FORMATS:
        return _err(f"Invalid model type '{model_type}'. Must be one of: {', '.join(sorted(VALID_FORMATS))}")

    try:
        files_list = json.loads(files)
        return _publish_product(
            project=project,
            folder_path=folder_path,
            product_name=product_name,
            product_type="modelMain",
            representation_name=model_type,
            files_list=files_list,
            variant=variant,
        )
    except Exception as e:
        return _err(str(e))


# ── Health ───────────────────────────────────────────────


@app.custom_route("/health", methods=["GET"])
async def health(request):  # noqa: ARG001
    from starlette.responses import JSONResponse

    return JSONResponse({"status": "ok"})


# ── Entry Point ──────────────────────────────────────────

def main():
    if not os.environ.get("AYON_URL"):
        print("ERROR: AYON_URL not set", file=sys.stderr)
        sys.exit(1)
    if not os.environ.get("AYON_KEY"):
        print("ERROR: AYON_KEY not set", file=sys.stderr)
        sys.exit(1)

    transport = os.environ.get("AYON_TRANSPORT", "stdio")
    if transport == "http":
        import uvicorn
        port = int(os.environ.get("PORT", "8000"))
        uvicorn.run(app.streamable_http_app(), host="0.0.0.0", port=port, log_level="info")
    else:
        app.run(transport="stdio")


if __name__ == "__main__":
    main()