server.py

"""PCSWMM Engineering MCP Server.

Local-first surfaces:
1. MCP Streamable HTTP at /mcp.
2. REST/OpenAPI at /api/*.
3. Optional narrative agent at /api/agent and agent_analyze.
4. Generated artifact download at /files/{session_id}/{filename}.

All public tools are PCSWMM-facing and dispatch through pcswmm_tools.py.
The generic uploaded-INP functions remain internal only for independent verification.
"""
from __future__ import annotations

import contextlib
import inspect
import json
import os
from pathlib import Path
from typing import Any

from fastapi import Body, FastAPI, HTTPException
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse

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

import agent as agent_module
from sessions import SESSION_ROOT, STORE
from pcswmm_tools import PCSWMM_TOOL_REGISTRY as TOOL_REGISTRY

SERVER_NAME = "pcswmm-engineering"
SERVER_VERSION = "1.10.0 (PySWMM model-path readiness fix; staged draft/review/submission gates; hydraulic summary)"

# ---------------------------------------------------------------------------
# MCP surface
# ---------------------------------------------------------------------------
# The session-manager lifespan is wired into FastAPI so Streamable HTTP remains active.
mcp = FastMCP(
    SERVER_NAME,
    instructions=(
        "PCSWMM-only engineering copilot. Connect either the active PCSWMM SDK package with "
        "connect_active_pcswmm_project or a prior deterministic evidence folder with connect_deterministic_calgary_evidence; use the declared deterministic backend (currently PCSWMM SDK or PySWMM/EPA SWMM), optionally verify the exact INP, review "
        "hydrology/hydraulics/storage and revision impacts, configure first or revised "
        "Calgary submissions, resolve City comments, and generate the SWMR and audit ZIP. "
        "The public workflow does not accept arbitrary standalone model uploads. All "
        "outputs are deterministic engineering screening for professional review."
    ),
    stateless_http=True,
    json_response=True,
    transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)

for _name, _fn in TOOL_REGISTRY.items():
    mcp.tool()(_fn)


@mcp.tool()
def agent_analyze(question: str, provider: str = "local", model: str = "",
                  session_id: str = "", api_key: str = "", base_url: str = "") -> dict:
    """Ask the built-in agent a natural-language question; it plans and runs the
    SWMM tools itself and returns an answer plus a full tool audit trail.
    Providers: local, anthropic, openai, gemini, groq, or mistral. The deterministic tools do not require an LLM. Use this endpoint only for optional narrative orchestration."""
    return agent_module.run_agent(
        question=question, provider=provider, model=model or None,
        api_key=api_key or None, base_url=base_url or None,
        session_id=session_id or None)


mcp_app = mcp.streamable_http_app()

# ---------------------------------------------------------------------------
# FastAPI app with the MCP lifespan wired in (fix #2)
# ---------------------------------------------------------------------------

@contextlib.asynccontextmanager
async def lifespan(app: FastAPI):
    async with contextlib.AsyncExitStack() as stack:
        await stack.enter_async_context(mcp.session_manager.run())
        yield


app = FastAPI(
    title="PCSWMM Engineering MCP Server",
    version=SERVER_VERSION,
    description="Local-first PCSWMM Engineering MCP. The service receives the active PCSWMM Engineering SDK package, uses the active PCSWMM evidence as the primary backend, optionally verifies the exact INP, performs deterministic QA/QC, Calgary screening, revision review, City-comment tracking, and SWMR generation.",
    lifespan=lifespan,
)


# ---------------------------------------------------------------------------
# REST surface — one endpoint per tool, plus a generic dispatcher
# ---------------------------------------------------------------------------

def _tool_meta(name: str, fn) -> dict:
    sig = inspect.signature(fn)
    return {
        "name": name,
        "description": (fn.__doc__ or "").strip(),
        "parameters": {
            p: {"required": prm.default is inspect.Parameter.empty,
                "default": None if prm.default is inspect.Parameter.empty else prm.default}
            for p, prm in sig.parameters.items()
        },
        "rest": f"POST /api/tool/{name}",
    }


@app.get("/api/tools")
def rest_list_tools() -> dict:
    """List all tools with parameter metadata (machine-readable)."""
    return {"server": SERVER_NAME, "version": SERVER_VERSION,
            "tools": [_tool_meta(n, f) for n, f in TOOL_REGISTRY.items()]}


@app.post("/api/tool/{tool_name}")
def rest_call_tool(tool_name: str, payload: dict = Body(default={})) -> JSONResponse:
    """Invoke any registry tool. Body = the tool's keyword arguments as JSON."""
    fn = TOOL_REGISTRY.get(tool_name)
    if fn is None:
        raise HTTPException(404, f"Unknown tool '{tool_name}'. See /api/tools.")
    try:
        result = fn(**(payload or {}))
    except (KeyError, ValueError, TypeError) as exc:
        raise HTTPException(400, str(exc))
    except Exception as exc:
        raise HTTPException(500, f"{type(exc).__name__}: {exc}")
    return JSONResponse(json.loads(json.dumps(result, default=str)))


@app.post("/api/agent")
def rest_agent(payload: dict = Body(...)) -> JSONResponse:
    """Built-in agent endpoint.

    Body: {"question": str, "provider": "anthropic|openai|gemini|groq|mistral|local",
           "model": str?, "api_key": str?, "base_url": str?, "session_id": str?,
           "inp_content": str?, "allow_report": bool?}
    """
    question = payload.get("question", "").strip()
    if not question:
        raise HTTPException(400, "'question' is required.")
    try:
        result = agent_module.run_agent(
            question=question,
            provider=payload.get("provider", "local"),
            model=payload.get("model") or None,
            api_key=payload.get("api_key") or None,
            base_url=payload.get("base_url") or None,
            session_id=payload.get("session_id") or None,
            inp_content=payload.get("inp_content") or None,
            allow_report=bool(payload.get("allow_report", False)),
        )
    except ValueError as exc:
        raise HTTPException(400, str(exc))
    except Exception as exc:
        raise HTTPException(502, f"Agent/provider error: {type(exc).__name__}: {exc}")
    return JSONResponse(json.loads(json.dumps(result, default=str)))


@app.get("/files/{session_id}/{filename}")
def serve_file(session_id: str, filename: str) -> FileResponse:
    """Download generated report artifacts."""
    safe_session = Path(session_id).name
    safe_file = Path(filename).name
    path = (SESSION_ROOT / safe_session / "outputs" / safe_file).resolve()
    if not str(path).startswith(str(SESSION_ROOT.resolve())) or not path.exists():
        raise HTTPException(404, "File not found (sessions expire; regenerate the report).")
    return FileResponse(path, filename=safe_file)


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "product": "PCSWMM Engineering MCP", "server": SERVER_NAME, "version": SERVER_VERSION, "mode": "local-first", "bind_recommendation": "127.0.0.1", "tools": len(TOOL_REGISTRY) + 1, "sessions": len(STORE.list()), "mcp_endpoint": "/mcp", "openapi": "/openapi.json"}


@app.get("/", response_class=HTMLResponse)
def index() -> str:
    tool_rows = "".join(
        f"<tr><td><code>{n}</code></td><td>{(f.__doc__ or '').strip().splitlines()[0]}</td></tr>"
        for n, f in TOOL_REGISTRY.items())
    return f"""<!doctype html><html><head><title>PCSWMM Engineering MCP</title>
<style>body{{font-family:system-ui;max-width:960px;margin:2rem auto;padding:0 1rem;color:#222}}
code{{background:#f2f2f2;padding:1px 5px;border-radius:4px}}table{{border-collapse:collapse;width:100%}}
td,th{{border:1px solid #ddd;padding:6px 10px;text-align:left;font-size:14px}}h1{{color:#0a4d8c}}</style></head>
<body><h1>PCSWMM Engineering MCP</h1>
<p>Local-first engineering copilot for consultants who use PCSWMM. The server starts from the active PCSWMM Engineering SDK package, uses the active PCSWMM evidence as the primary backend, optionally performs independent SWMM verification, and prepares engineering review and Calgary SWMR deliverables.</p>
<ul><li><b>MCP:</b> <code>http://127.0.0.1:8765/mcp</code></li>
<li><b>REST catalog:</b> <a href="/api/tools">/api/tools</a></li>
<li><b>Health:</b> <a href="/health">/health</a></li></ul>
<p><b>Required first tool:</b> <code>connect_active_pcswmm_project</code>. Standalone uploaded-INP workflow is intentionally not exposed.</p>
<h3>PCSWMM tools ({len(TOOL_REGISTRY) + 1})</h3><table><tr><th>Tool</th><th>Purpose</th></tr>{tool_rows}
<tr><td><code>agent_analyze</code></td><td>Optional local/cloud narrative orchestration over the PCSWMM tools.</td></tr></table></body></html>"""

# Mount MCP last so explicit routes take priority.
app.mount("/", mcp_app)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=os.environ.get("PCSWMM_MCP_HOST", "127.0.0.1"), port=int(os.environ.get("PCSWMM_MCP_PORT", "8765")))