server/web.py

#!/usr/bin/env python3
"""
NeuralCAD Web Demo Server
=========================
FastAPI server that proxies REST requests to the MCP CAD server (SSE transport)
and serves the web frontend.

Usage:
    # Start MCP server first:
    python -m server.mcp --transport sse --port 8000

    # Then start web server:
    python -m server.web

    # Or auto-launch MCP server:
    python -m server.web --start-mcp

    # Open http://localhost:5000
"""

import json
import os
import subprocess
import sys
import tempfile
import time
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, File, Form, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse

from server.routes import router

from mcp import ClientSession
from mcp.client.sse import sse_client

# ── Config ───────────────────────────────────────────────────────────────

from config.settings import settings

OUTPUT_DIR = settings.output_dir
if not OUTPUT_DIR.is_absolute():
    OUTPUT_DIR = Path(__file__).parent.parent / OUTPUT_DIR

WEB_DIR = Path(settings.paths.web_dir)
if not WEB_DIR.is_absolute():
    WEB_DIR = Path(__file__).parent.parent / WEB_DIR

PORT = settings.web_port

MCP_SERVER_URL = os.environ.get("MCP_SERVER_URL", f"http://localhost:{settings.mcp_port}/sse")

# ── MCP Client Management ───────────────────────────────────────────────

_mcp_process = None


async def call_mcp_tool(tool_name: str, arguments: dict) -> dict:
    """Connect to MCP server, call a tool, return parsed JSON result."""
    async with sse_client(url=MCP_SERVER_URL) as streams:
        async with ClientSession(*streams) as session:
            await session.initialize()
            result = await session.call_tool(name=tool_name, arguments=arguments)
            if result.content:
                return json.loads(result.content[0].text)
            return {"error": "Empty response from MCP server"}


async def read_mcp_resource(uri: str) -> str:
    """Connect to MCP server and read a resource."""
    async with sse_client(url=MCP_SERVER_URL) as streams:
        async with ClientSession(*streams) as session:
            await session.initialize()
            result = await session.read_resource(uri=uri)
            if result.contents:
                return result.contents[0].text
            return "{}"


def start_mcp_server(port: int = 8000):
    """Launch mcp.py as a subprocess with SSE transport."""
    global _mcp_process
    mcp_script = Path(__file__).parent / "mcp.py"
    _mcp_process = subprocess.Popen(
        [sys.executable, str(mcp_script), "--transport", "sse", "--port", str(port)],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    # Give it a moment to start
    time.sleep(2)
    if _mcp_process.poll() is not None:
        stderr = _mcp_process.stderr.read().decode() if _mcp_process.stderr else ""
        raise RuntimeError(f"MCP server failed to start: {stderr}")
    print(f"  MCP server started (PID {_mcp_process.pid}) on port {port}")


# ── FastAPI App ──────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    OUTPUT_DIR.mkdir(exist_ok=True)
    yield
    global _mcp_process
    if _mcp_process:
        _mcp_process.terminate()
        _mcp_process.wait()


app = FastAPI(title="NeuralCAD Web Demo", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(router)


# ── Routes ───────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def index():
    index_file = WEB_DIR / "index.html"
    return HTMLResponse(index_file.read_text())


@app.post("/api/generate")
async def generate(body: dict):
    result = await call_mcp_tool("generate_cnc_model", {
        "prompt": body.get("prompt", ""),
        "part_name": body.get("part_name", ""),
        "backend": body.get("backend", "mock"),
        "max_retries": body.get("max_retries", 2),
    })
    return JSONResponse(result)


@app.post("/api/generate-image")
async def generate_image(
    image: UploadFile = File(...),
    text_hint: str = Form(""),
    part_name: str = Form(""),
    backend: str = Form("anthropic"),
):
    # Save uploaded image to temp file
    suffix = Path(image.filename or "upload.png").suffix
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(await image.read())
        tmp_path = tmp.name

    try:
        result = await call_mcp_tool("generate_from_image", {
            "image_path": tmp_path,
            "text_hint": text_hint,
            "part_name": part_name,
            "backend": backend,
        })
        return JSONResponse(result)
    finally:
        os.unlink(tmp_path)


@app.post("/api/validate")
async def validate(body: dict):
    result = await call_mcp_tool("validate_cnc_model", {
        "cadquery_code": body.get("code", ""),
        "part_name": body.get("part_name", "Part"),
    })
    return JSONResponse(result)


@app.get("/api/models")
async def list_models():
    result = await call_mcp_tool("list_models", {
        "output_dir": str(OUTPUT_DIR),
    })
    return JSONResponse(result)


import re

_SAFE_NAME = re.compile(r'^[a-zA-Z0-9_\-]+$')


def _safe_model_path(name: str, ext: str) -> Path | None:
    """Validate model name and return safe path, or None if invalid."""
    if not _SAFE_NAME.match(name):
        return None
    path = (OUTPUT_DIR / f"{name}.{ext}").resolve()
    if not str(path).startswith(str(OUTPUT_DIR.resolve())):
        return None
    return path


@app.get("/api/models/{name}.stl")
async def get_stl(name: str):
    path = _safe_model_path(name, "stl")
    if not path or not path.exists():
        return JSONResponse({"error": f"STL not found: {name}"}, status_code=404)
    return FileResponse(path, media_type="model/stl", filename=f"{name}.stl")


@app.get("/api/models/{name}.step")
async def get_step(name: str):
    path = _safe_model_path(name, "step")
    if not path or not path.exists():
        return JSONResponse({"error": f"STEP not found: {name}"}, status_code=404)
    return FileResponse(path, media_type="application/step", filename=f"{name}.step")


@app.get("/api/models/{name}.gcode")
async def get_gcode(name: str):
    path = _safe_model_path(name, "gcode")
    if not path or not path.exists():
        return JSONResponse({"error": f"G-code not found: {name}"}, status_code=404)
    return FileResponse(path, media_type="text/plain", filename=f"{name}.gcode")


@app.get("/api/models/{name}.3mf")
async def get_3mf(name: str):
    path = _safe_model_path(name, "3mf")
    if not path or not path.exists():
        return JSONResponse({"error": f"3MF not found: {name}"}, status_code=404)
    return FileResponse(path, media_type="model/3mf", filename=f"{name}.3mf")


@app.get("/api/capabilities")
async def capabilities():
    try:
        text = await read_mcp_resource("text-to-cnc://capabilities")
        return JSONResponse(json.loads(text))
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=502)


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

if __name__ == "__main__":
    import argparse
    import uvicorn

    parser = argparse.ArgumentParser(description="NeuralCAD Web Demo Server")
    parser.add_argument("--port", type=int, default=PORT, help="Web server port (default: 5000)")
    parser.add_argument("--host", default="0.0.0.0", help="Bind host (default: 0.0.0.0)")
    parser.add_argument(
        "--start-mcp", action="store_true",
        help="Auto-launch MCP server as subprocess before starting web server"
    )
    parser.add_argument("--mcp-port", type=int, default=8000, help="MCP server port (default: 8000)")
    args = parser.parse_args()

    if args.start_mcp:
        MCP_SERVER_URL = f"http://localhost:{args.mcp_port}/sse"
        print(f"Starting MCP CAD server on port {args.mcp_port}...")
        start_mcp_server(args.mcp_port)

    print(f"Starting NeuralCAD Web Demo on http://localhost:{args.port}")
    print(f"MCP server: {MCP_SERVER_URL}")
    uvicorn.run(app, host=args.host, port=args.port)