quantux_mcp/server.py

"""
Quant-UX MCP Server — exposes the Quant-UX prototype tool to any MCP client
(Claude Code, Codex CLI, Cursor, ...) over Streamable HTTP.

Environment variables:
  QUX_BASE_URL       Quant-UX frontend base URL (default http://127.0.0.1:8082)
  QUX_ADMIN_EMAIL    optional account used to auto-login at startup
  QUX_ADMIN_PASSWORD optional password for the account above
  MCP_API_KEY        optional Bearer key required on every HTTP request
                     (strongly recommended when exposed on a public server)
  MCP_PORT           HTTP port (default 8090)

Run:  python -m uvicorn server:app --host 0.0.0.0 --port 8090
"""

import json
import os

from mcp.server.mcpserver import MCPServer

from quantux_client import QuantUXClient, QuantUXError
from quantux_auth import DownloadTokenStore, QuantUXSessionManager
from quantux_export import QuantUXExporter
from quantux_verify import verify_html

BASE_URL = os.environ.get("QUX_BASE_URL", "http://127.0.0.1:8082")
ADMIN_EMAIL = os.environ.get("QUX_ADMIN_EMAIL", "")
ADMIN_PASSWORD = os.environ.get("QUX_ADMIN_PASSWORD", "")
EXPORT_DIR = os.environ.get("QUX_EXPORT_DIR", "/app/exports")
os.makedirs(EXPORT_DIR, exist_ok=True)

_client = QuantUXClient(BASE_URL)
_session = QuantUXSessionManager(_client, ADMIN_EMAIL, ADMIN_PASSWORD)
_client.set_auth_refresh_handler(_session.refresh_after_auth_failure)


def _auto_login():
    """Log in at startup if admin credentials are configured."""
    if ADMIN_EMAIL and ADMIN_PASSWORD:
        try:
            _session.ensure_authenticated(force=True)
            print(f"Auto-login ok as {ADMIN_EMAIL}", flush=True)
        except Exception as exc:  # noqa: BLE001
            print(f"WARNING: auto-login failed: {exc}", flush=True)


_auto_login()


def _client_or_error():
    return _session.ensure_authenticated()


def _ok(data):
    if isinstance(data, (dict, list)):
        return json.dumps(data, ensure_ascii=False)
    return str(data)


server = MCPServer(
    "quantux",
    title="Quant-UX Design Server",
    description=(
        "Create, edit and wire interactive UI prototypes on a self-hosted "
        "Quant-UX instance. Tools cover accounts, apps, screens, widgets "
        "(Box/Label/Button/TextBox/Password/Image/...), styles and flow "
        "connections."
    ),
    version="1.0.0",
)


# ---------------------------------------------------------------- accounts
@server.tool(
    name="quantux_login",
    title="Login",
    description="Log into Quant-UX with an email/password. Required before any "
                "other tool if the server has no admin account configured.",
)
def quantux_login(email: str, password: str) -> str:
    """Authenticate against Quant-UX and cache the JWT."""
    user = _client.login(email, password)
    return _ok({
        "status": "ok",
        "email": user.get("email"),
        "name": user.get("name"),
        "role": user.get("role"),
    })


@server.tool(
    name="quantux_register",
    title="Register account",
    description="Create a new Quant-UX user account.",
)
def quantux_register(name: str, lastname: str, email: str, password: str) -> str:
    user = _client.register(name, lastname, email, password)
    return _ok({"status": "ok", "id": user.get("_id"), "email": user.get("email")})


# ------------------------------------------------------------------- apps
@server.tool(
    name="quantux_list_apps",
    title="List apps",
    description="List all prototypes/apps owned by the logged-in user.",
)
def quantux_list_apps() -> str:
    apps = _client_or_error().list_apps()
    summary = [{
        "id": a.get("_id") or a.get("id"),
        "name": a.get("name"),
        "type": a.get("type"),
        "screenSize": a.get("screenSize"),
        "isPublic": a.get("isPublic", False),
    } for a in apps]
    return _ok(summary)


@server.tool(
    name="quantux_get_app",
    title="Inspect app",
    description="Return a structured summary of an app: screens, widget counts "
                "and flow lines. Use quantux_dump_app for the raw model.",
)
def quantux_get_app(app_id: str) -> str:
    return _ok(_client_or_error().describe(app_id))


@server.tool(
    name="quantux_dump_app",
    title="Dump raw app model",
    description="Return the full raw Quant-UX model JSON of an app "
                "(screens/widgets/lines with all styles).",
)
def quantux_dump_app(app_id: str) -> str:
    return _ok(_client_or_error().get_app(app_id))


@server.tool(
    name="quantux_create_app",
    title="Create app",
    description="Create a new prototype. Defaults to a 375x667 smartphone "
                "canvas; pass width/height for desktop (e.g. 1280x720).",
)
def quantux_create_app(
    name: str,
    description: str = "",
    width: int = 375,
    height: int = 667,
    app_type: str = "prototype",
) -> str:
    app_id = _client_or_error().create_app(
        name, description=description, width=width, height=height,
        app_type=app_type,
    )
    return _ok({"status": "ok", "app_id": app_id, "name": name})


@server.tool(
    name="quantux_delete_app",
    title="Delete app",
    description="Permanently delete an app.",
)
def quantux_delete_app(app_id: str) -> str:
    _client_or_error().delete_app(app_id)
    return _ok({"status": "ok", "deleted": app_id})


# ---------------------------------------------------------------- screens
@server.tool(
    name="quantux_add_screen",
    title="Add screen",
    description="Add a screen to an app. The first screen automatically "
                "becomes the start screen.",
)
def quantux_add_screen(app_id: str, name: str, width: int = 0, height: int = 0) -> str:
    c = _client_or_error()
    sid = c.add_screen(
        app_id, name,
        width=width or None, height=height or None,
    )
    return _ok({"status": "ok", "screen_id": sid, "name": name})


# ---------------------------------------------------------------- widgets
@server.tool(
    name="quantux_add_widget",
    title="Add widget",
    description=(
        "Add a widget to a screen. widget_type is one of: Box (rectangle), "
        "Label (text), Button, TextBox (input), Password, TextArea, Image, "
        "Icon, HotSpot. Coordinates x/y are absolute within the screen; "
        "w/h are width/height. Style keys are CSS-ish: background, color, "
        "fontSize, fontWeight, textAlign, borderRadius corners, border*Width, "
        "padding*, boxShadow, lineHeight, letterSpacing, fontFamily."
    ),
)
def quantux_add_widget(
    app_id: str,
    screen_id: str,
    widget_type: str,
    x: int,
    y: int,
    w: int,
    h: int,
    name: str = "",
    props_json: str = "{}",
    style_json: str = "{}",
) -> str:
    props = json.loads(props_json or "{}")
    style = json.loads(style_json or "{}")
    wid = _client_or_error().add_widget(
        app_id, screen_id, widget_type, x, y, w, h,
        name=name or None, props=props, style=style,
    )
    return _ok({"status": "ok", "widget_id": wid, "type": widget_type})


@server.tool(
    name="quantux_update_widget",
    title="Update widget",
    description="Patch a widget: style_json/props_json are merged onto the "
                "existing values; x/y/w/h/name replace position/size/name.",
)
def quantux_update_widget(
    app_id: str,
    widget_id: str,
    style_json: str = "{}",
    props_json: str = "{}",
    x: int = -1,
    y: int = -1,
    w: int = -1,
    h: int = -1,
    name: str = "",
) -> str:
    c = _client_or_error()
    c.update_widget(
        app_id, widget_id,
        props=json.loads(props_json or "{}") or None,
        style=json.loads(style_json or "{}") or None,
        x=x if x >= 0 else None,
        y=y if y >= 0 else None,
        w=w if w >= 0 else None,
        h=h if h >= 0 else None,
        name=name or None,
    )
    return _ok({"status": "ok", "widget_id": widget_id})


@server.tool(
    name="quantux_delete_widget",
    title="Delete widget",
    description="Remove a widget from its screen.",
)
def quantux_delete_widget(app_id: str, widget_id: str) -> str:
    _client_or_error().delete_widget(app_id, widget_id)
    return _ok({"status": "ok", "deleted": widget_id})


# ------------------------------------------------------------------- flows
@server.tool(
    name="quantux_connect_flow",
    title="Connect flow",
    description=(
        "Wire an interaction: clicking the 'from' widget navigates to the "
        "'to' screen (or widget). event defaults to 'click'."
    ),
)
def quantux_connect_flow(
    app_id: str,
    from_widget_id: str,
    to_screen_id: str,
    event: str = "click",
) -> str:
    _client_or_error().connect_flow(app_id, from_widget_id, to_screen_id, event)
    return _ok({"status": "ok", "from": from_widget_id,
                "to": to_screen_id, "event": event})


# ------------------------------------------------------------- escape hatch
@server.tool(
    name="quantux_apply_changes",
    title="Apply raw changes",
    description=(
        "Low-level escape hatch: apply a raw Quant-UX delta array to an app. "
        "Each change: {\"type\":\"add|update|delete\", \"name\":<id or field>, "
        "\"parent\":\"screens|widgets|lines|groups|null\", \"object\":<value>}. "
        "Use for advanced edits not covered by the other tools."
    ),
)
def quantux_apply_changes(app_id: str, changes_json: str) -> str:
    changes = json.loads(changes_json)
    _client_or_error().apply_changes(app_id, changes)
    return _ok({"status": "ok", "applied": len(changes)})


# ------------------------------------------------------------------- health
@server.tool(
    name="quantux_export_html",
    title="Export interactive HTML",
    description=(
        "Export a prototype to a single self-contained interactive HTML file "
        "that works fully offline (open in any browser, click through flows). "
        "Images are embedded as base64. The export is automatically verified "
        "for interaction consistency; the 'verify' field reports the checks. "
        "The result contains a short-lived download URL scoped to that file "
        "(or pull from the server's ~/quantux-mcp/exports/ directory)."
    ),
)
def quantux_export_html(app_id: str) -> str:
    c = _client_or_error()
    path = QuantUXExporter(c).export(app_id, out_path=os.path.join(EXPORT_DIR, f"{app_id}.html"))
    try:
        with open(path, encoding="utf-8") as fh:
            verify = verify_html(fh.read())
    except Exception as exc:  # noqa: BLE001
        verify = {"status": "error", "error": str(exc)[:200]}
    filename = os.path.basename(path)
    download_token = _download_tokens.issue(filename)
    return _ok({
        "status": "ok",
        "app_id": app_id,
        "file": filename,
        "bytes": os.path.getsize(path),
        "download": f"/exports/{filename}?download_token={download_token}",
        "download_expires_in_seconds": _download_tokens.ttl_seconds,
        "verify": {
            "verdict": verify.get("verdict"),
            "widgetCount": verify.get("widgetCount"),
            "wiredCount": verify.get("wiredCount"),
            "navigation": verify.get("navigation"),
            "inputsEditable": verify.get("inputsEditable"),
            "toggleWorks": verify.get("toggleWorks"),
            "charts": verify.get("charts", 0),
            "tables": verify.get("tables", 0),
            "repeater": None,
            "animation": verify.get("animation"),
            "runtimeErrors": len(verify.get("runtimeErrors") or []),
            "unsupportedTypes": list((verify.get("model") or {}).get("unsupported", {}).keys()),
        },
    })


@server.tool(
    name="quantux_verify_export",
    title="Verify export consistency",
    description=(
        "Export a prototype and run the full consistency verification "
        "(render, wiring, navigation, inputs, toggles, charts/tables, "
        "animation, widget-type coverage). Returns the detailed report."
    ),
)
def quantux_verify_export(app_id: str) -> str:
    c = _client_or_error()
    path = QuantUXExporter(c).export(app_id, out_path=os.path.join(EXPORT_DIR, f"{app_id}.html"))
    with open(path, encoding="utf-8") as fh:
        report = verify_html(fh.read())
    report["file"] = os.path.basename(path)
    return _ok(report)


@server.tool(
    name="quantux_health",
    title="Health check",
    description="Check MCP server status and Quant-UX backend reachability.",
)
def quantux_health() -> str:
    status = {"mcp": "ok", "quantux_base": BASE_URL}
    auth_error = None
    try:
        if ADMIN_EMAIL and ADMIN_PASSWORD:
            _session.ensure_authenticated()
    except Exception as exc:  # noqa: BLE001
        auth_error = str(exc)[:200]
    try:
        r = _client._req("GET", "/rest/status.json")
        status["quantux"] = "ok"
        status["backend"] = r.get("version")
    except Exception as exc:  # noqa: BLE001
        status["quantux"] = f"error: {exc}"
    status.update(_session.status())
    if auth_error:
        status["auth_error"] = auth_error
    return _ok(status)


# ------------------------------------------------------------------- app
# Auth middleware: require Bearer MCP_API_KEY on every request.
_API_KEY = os.environ.get("MCP_API_KEY", "")
_download_tokens = DownloadTokenStore(
    ttl_seconds=int(os.environ.get("MCP_EXPORT_TOKEN_TTL_SECONDS", "600"))
)


def _auth_required(request):
    if not _API_KEY:
        return None
    auth = request.headers.get("authorization", "")
    expected = f"Bearer {_API_KEY}"
    if auth != expected:
        return {"error": "unauthorized", "detail": "invalid or missing API key"}
    return None


def make_app():
    from starlette.middleware.base import BaseHTTPMiddleware
    from starlette.responses import JSONResponse
    from starlette.staticfiles import StaticFiles

    starlette_app = server.streamable_http_app(
        streamable_http_path="/mcp",
        # Explicit non-localhost host disables the SDK's automatic DNS
        # rebinding protection (which would otherwise 421-reject requests
        # whose Host header is not 127.0.0.1). Access control is provided
        # by MCP_API_KEY auth plus the cloud security group on the port.
        host=os.environ.get("MCP_HOST", "0.0.0.0"),
    )

    # Serve exported prototypes at /exports/<file>?key=<MCP_API_KEY>
    # (GET with the key as a query param so it works in a plain browser).
    starlette_app.mount(
        "/exports",
        StaticFiles(directory=EXPORT_DIR),
        name="exports",
    )

    class AuthMiddleware(BaseHTTPMiddleware):
        async def dispatch(self, request, call_next):
            problem = _auth_required(request)
            if problem:
                if request.method == "GET" and request.url.path.startswith("/exports/"):
                    filename = os.path.basename(request.url.path)
                    token = request.query_params.get("download_token", "")
                    if _download_tokens.validate(token, filename):
                        return await call_next(request)
                return JSONResponse(problem, status_code=401)
            return await call_next(request)

    starlette_app.add_middleware(AuthMiddleware)
    return starlette_app


app = make_app()

if __name__ == "__main__":
    import uvicorn

    port = int(os.environ.get("MCP_PORT", "8090"))
    uvicorn.run(app, host="0.0.0.0", port=port)