server.py

"""
EPANET MCP Server — Hugging Face Spaces deployment
Uses the official `mcp` Python SDK (FastMCP + SSE transport).
Compatible with Claude Desktop (mcp-remote), claude.ai connectors, and any MCP client.
"""

import asyncio
import json
import uuid
import tempfile
import os
from typing import Any

from mcp.server.fastmcp import FastMCP
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Route, Mount
import uvicorn

# ── Client API key (set as an HF Space secret named CLIENT_API_KEY) ────────────
CLIENT_API_KEY = os.environ.get("CLIENT_API_KEY")  # required to use the MCP tools

# Paths that stay open without a key
_PUBLIC_PATHS = {"/", "/health"}

_UNAUTHORIZED = JSONResponse(
    {"error": "Unauthorized. Missing or invalid API key."}, status_code=401
).body

_MISCONFIGURED = JSONResponse(
    {"error": "Server misconfigured: CLIENT_API_KEY not set."}, status_code=500
).body


class APIKeyMiddleware:
    """Pure ASGI middleware — compatible with streaming responses (SSE + Streamable HTTP)."""

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] not in ("http", "websocket"):
            await self.app(scope, receive, send)
            return

        path = scope.get("path", "")
        if path in _PUBLIC_PATHS:
            await self.app(scope, receive, send)
            return

        # Extract Authorization header
        headers = dict(scope.get("headers", []))
        auth = headers.get(b"authorization", b"").decode()
        provided = auth.removeprefix("Bearer ").strip()

        if not CLIENT_API_KEY:
            await self._reject(scope, send, 500, _MISCONFIGURED)
            return

        if provided != CLIENT_API_KEY:
            await self._reject(scope, send, 401, _UNAUTHORIZED)
            return

        await self.app(scope, receive, send)

    @staticmethod
    async def _reject(scope, send, status: int, body: bytes):
        await send({
            "type": "http.response.start",
            "status": status,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
            ],
        })
        await send({"type": "http.response.body", "body": body})

# ── in-memory session store ────────────────────────────────────────────────────
_sessions: dict[str, Any] = {}   # network_id → epyt EPANET object
_tmp_files: dict[str, str] = {}  # network_id → temp .inp path

# ── FastMCP instance ───────────────────────────────────────────────────────────
mcp = FastMCP("EPANET MCP Server")


# ── helpers ────────────────────────────────────────────────────────────────────

def _get_session(network_id: str):
    if network_id not in _sessions:
        raise ValueError(f"Network '{network_id}' not loaded. Call load_network first.")
    return _sessions[network_id]


def _bundled_path(filename: str) -> str:
    import epyt
    base = os.path.dirname(epyt.__file__)
    networks_dir = os.path.join(base, "networks")
    for root, _, files in os.walk(networks_dir):
        for f in files:
            if f.lower() == filename.lower():
                return os.path.join(root, f)
    return filename


def _network_summary(d) -> dict:
    return {
        "nodes": {
            "total": d.getNodeCount(),
            "junctions": d.getNodeJunctionCount(),
            "tanks": d.getNodeTankCount(),
            "reservoirs": d.getNodeReservoirCount(),
        },
        "links": {
            "total": d.getLinkCount(),
            "pipes": d.getLinkPipeCount(),
            "pumps": d.getLinkPumpCount(),
            "valves": d.getLinkValveCount(),
        },
        "patterns": d.getPatternCount(),
        "curves": d.getCurveCount(),
        "simple_controls": d.getControlRulesCount(),
        "flow_units": d.getFlowUnits(),
        "simulation_duration_s": int(d.getTimeSimulationDuration()),
        "hydraulic_timestep_s": int(d.getTimeHydraulicStep()),
    }


# ── MCP tools ──────────────────────────────────────────────────────────────────

@mcp.tool()
def list_bundled_networks() -> dict:
    """List all .inp benchmark networks bundled with ePyT."""
    import epyt
    base = os.path.dirname(epyt.__file__)
    networks_dir = os.path.join(base, "networks")
    results = []
    for root, _, files in os.walk(networks_dir):
        for f in files:
            if f.endswith(".inp") and "_temp" not in f:
                rel = os.path.relpath(os.path.join(root, f), networks_dir)
                results.append(rel)
    return {"bundled_networks": sorted(results)}


@mcp.tool()
def list_networks() -> dict:
    """Return IDs of all currently loaded network sessions."""
    return {"network_ids": list(_sessions.keys())}


@mcp.tool()
def load_network(path: str, network_id: str = None) -> dict:
    """
    Load an EPANET .inp file into a named session.
    path: bundled filename (e.g. 'Net1.inp') or absolute path.
    network_id: optional session name; defaults to base filename.
    """
    from epyt import epanet as EPANET
    nid = network_id or os.path.splitext(os.path.basename(path))[0]
    resolved = _bundled_path(path) if not os.path.isabs(path) else path
    d = EPANET(resolved, verbose=False)
    _sessions[nid] = d
    return {"network_id": nid, "path": resolved, "summary": _network_summary(d)}


@mcp.tool()
def unload_network(network_id: str) -> dict:
    """Unload a network session and free its resources."""
    if network_id in _sessions:
        try:
            _sessions[network_id].unload()
        except Exception:
            pass
        del _sessions[network_id]
    if network_id in _tmp_files:
        try:
            os.unlink(_tmp_files.pop(network_id))
        except Exception:
            pass
    return {"unloaded": True, "network_id": network_id}


@mcp.tool()
def get_network_summary(network_id: str) -> dict:
    """Return node/link counts, flow units, and simulation settings."""
    return _network_summary(_get_session(network_id))


@mcp.tool()
def get_nodes(network_id: str) -> dict:
    """Return all nodes (junctions, tanks, reservoirs) with elevation, demand, and type."""
    d = _get_session(network_id)
    ids = d.getNodeNameID()
    elevs = d.getNodeElevations()
    demands = d.getNodeBaseDemands()[1]
    types = d.getNodeType()
    nodes = []
    for i, nid in enumerate(ids):
        nodes.append({
            "id": nid,
            "type": types[i],
            "elevation": float(elevs[i]),
            "base_demand": float(demands[i]),
        })
    return {"nodes": nodes}


@mcp.tool()
def get_links(network_id: str) -> dict:
    """Return all links (pipes, pumps, valves) with diameter, length, roughness, and status."""
    d = _get_session(network_id)
    ids = d.getLinkNameID()
    types = d.getLinkType()
    diams = d.getLinkDiameter()
    lengths = d.getLinkLength()
    rough = d.getLinkRoughnessCoeff()
    nodes_from = d.getLinkNodesIndex()
    node_names = d.getNodeNameID()
    links = []
    for i, lid in enumerate(ids):
        nf = nodes_from[0][i] - 1
        nt = nodes_from[1][i] - 1
        links.append({
            "id": lid,
            "type": types[i],
            "from_node": node_names[nf],
            "to_node": node_names[nt],
            "diameter": float(diams[i]),
            "length": float(lengths[i]),
            "roughness": float(rough[i]),
        })
    return {"links": links}


@mcp.tool()
def run_hydraulic_simulation(network_id: str) -> dict:
    """Run a full hydraulic EPS simulation. Returns pressures, flows, velocities at every timestep."""
    d = _get_session(network_id)
    res = d.getComputedHydraulicTimeSeries()
    times = [int(t) for t in res.Time]
    node_ids = d.getNodeNameID()
    link_ids = d.getLinkNameID()
    pressures = {node_ids[i]: [round(float(v), 3) for v in res.Pressure[:, i]]
                 for i in range(len(node_ids))}
    flows = {link_ids[i]: [round(float(v), 3) for v in res.Flow[:, i]]
             for i in range(len(link_ids))}
    return {"times_s": times, "pressures": pressures, "flows": flows}


@mcp.tool()
def run_full_simulation(network_id: str) -> dict:
    """Run hydraulic + water quality simulation."""
    d = _get_session(network_id)
    res = d.getComputedTimeSeries()
    times = [int(t) for t in res.Time]
    node_ids = d.getNodeNameID()
    link_ids = d.getLinkNameID()
    pressures = {node_ids[i]: [round(float(v), 3) for v in res.Pressure[:, i]]
                 for i in range(len(node_ids))}
    flows = {link_ids[i]: [round(float(v), 3) for v in res.Flow[:, i]]
             for i in range(len(link_ids))}
    quality = {node_ids[i]: [round(float(v), 4) for v in res.NodeQuality[:, i]]
               for i in range(len(node_ids))}
    return {"times_s": times, "pressures": pressures, "flows": flows, "node_quality": quality}


@mcp.tool()
def get_pressure_at_time(network_id: str, time_s: int) -> dict:
    """Return node pressures at a specific simulation time (seconds)."""
    d = _get_session(network_id)
    res = d.getComputedHydraulicTimeSeries()
    times = [int(x) for x in res.Time]
    idx = min(range(len(times)), key=lambda i: abs(times[i] - time_s))
    node_ids = d.getNodeNameID()
    return {
        "time_s": times[idx],
        "pressures": {node_ids[i]: round(float(res.Pressure[idx, i]), 3)
                      for i in range(len(node_ids))},
    }


@mcp.tool()
def get_flow_at_time(network_id: str, time_s: int) -> dict:
    """Return link flows at a specific simulation time (seconds)."""
    d = _get_session(network_id)
    res = d.getComputedHydraulicTimeSeries()
    times = [int(x) for x in res.Time]
    idx = min(range(len(times)), key=lambda i: abs(times[i] - time_s))
    link_ids = d.getLinkNameID()
    return {
        "time_s": times[idx],
        "flows": {link_ids[i]: round(float(res.Flow[idx, i]), 3)
                  for i in range(len(link_ids))},
    }


@mcp.tool()
def get_controls(network_id: str) -> dict:
    """List all simple controls defined in the network."""
    d = _get_session(network_id)
    count = d.getControlRulesCount()
    controls = []
    for i in range(1, count + 1):
        ctrl = d.getControls(i)
        controls.append({"index": i, "control": str(ctrl)})
    return {"controls": controls}


@mcp.tool()
def get_patterns(network_id: str) -> dict:
    """Return all demand/operational patterns."""
    d = _get_session(network_id)
    ids = d.getPatternNameID()
    patterns = []
    for i, pid in enumerate(ids):
        vals = d.getPattern(i + 1).tolist()
        patterns.append({"id": pid, "values": [round(v, 4) for v in vals]})
    return {"patterns": patterns}


@mcp.tool()
def set_pipe_diameter(network_id: str, pipe_id: str, diameter: float) -> dict:
    """Set the diameter of a pipe (mm or inches depending on network units)."""
    d = _get_session(network_id)
    idx = d.getLinkIndex(pipe_id)
    d.setLinkDiameter(idx, diameter)
    return {"pipe_id": pipe_id, "diameter": diameter}


@mcp.tool()
def set_pipe_roughness(network_id: str, pipe_id: str, roughness: float) -> dict:
    """Set the Hazen-Williams C factor (or D-W roughness) for a pipe."""
    d = _get_session(network_id)
    idx = d.getLinkIndex(pipe_id)
    d.setLinkRoughnessCoeff(idx, roughness)
    return {"pipe_id": pipe_id, "roughness": roughness}


@mcp.tool()
def set_pipe_status(network_id: str, pipe_id: str, status: str) -> dict:
    """Open or close a pipe. status must be 'OPEN' or 'CLOSED'."""
    d = _get_session(network_id)
    idx = d.getLinkIndex(pipe_id)
    d.setLinkInitialStatus(idx, 1 if status == "OPEN" else 0)
    return {"pipe_id": pipe_id, "status": status}


@mcp.tool()
def set_pump_status(network_id: str, pump_id: str, status: str) -> dict:
    """Start (OPEN) or stop (CLOSED) a pump."""
    d = _get_session(network_id)
    idx = d.getLinkIndex(pump_id)
    d.setLinkInitialStatus(idx, 1 if status == "OPEN" else 0)
    return {"pump_id": pump_id, "status": status}


@mcp.tool()
def set_node_base_demand(network_id: str, node_id: str, demand: float) -> dict:
    """Set the base demand for a junction in the network's flow units."""
    d = _get_session(network_id)
    idx = d.getNodeIndex(node_id)
    d.setNodeBaseDemands(idx, demand)
    return {"node_id": node_id, "demand": demand}


@mcp.tool()
def set_reservoir_head(network_id: str, reservoir_id: str, head: float) -> dict:
    """Set the total head of a reservoir node."""
    d = _get_session(network_id)
    idx = d.getNodeIndex(reservoir_id)
    d.setNodeElevations(idx, head)
    return {"reservoir_id": reservoir_id, "head": head}


@mcp.tool()
def create_leakage_event(network_id: str, pipe_id: str, leak_fraction: float,
                         scenario_id: str = None) -> dict:
    """Simulate a pipe burst by adding a leakage emitter at the pipe's downstream node."""
    import epyt
    src = _get_session(network_id)
    sid = scenario_id or f"leak_{pipe_id}_{uuid.uuid4().hex[:6]}"
    tmp = tempfile.NamedTemporaryFile(suffix=".inp", delete=False)
    tmp.close()
    src.saveInputFile(tmp.name)
    d2 = epyt.epanet(tmp.name, verbose=False)
    _sessions[sid] = d2
    _tmp_files[sid] = tmp.name
    pipe_idx = d2.getLinkIndex(pipe_id)
    to_node = d2.getLinkNodesIndex()[1][pipe_idx - 1]
    d2.setNodeEmitterCoeff(to_node, leak_fraction * 100)
    res = d2.getComputedHydraulicTimeSeries()
    times = [int(t) for t in res.Time]
    node_ids = d2.getNodeNameID()
    pressures = {node_ids[i]: [round(float(v), 3) for v in res.Pressure[:, i]]
                 for i in range(len(node_ids))}
    return {"scenario_id": sid, "pipe_id": pipe_id, "leak_fraction": leak_fraction,
            "times_s": times, "pressures": pressures}


@mcp.tool()
def create_contamination_event(network_id: str, source_node_id: str, concentration: float,
                               start_time_s: int, end_time_s: int,
                               source_type: str = "CONCEN", scenario_id: str = None) -> dict:
    """Inject a contaminant at a node and simulate water quality propagation."""
    import epyt
    src = _get_session(network_id)
    sid = scenario_id or f"contam_{source_node_id}_{uuid.uuid4().hex[:6]}"
    tmp = tempfile.NamedTemporaryFile(suffix=".inp", delete=False)
    tmp.close()
    src.saveInputFile(tmp.name)
    d2 = epyt.epanet(tmp.name, verbose=False)
    _sessions[sid] = d2
    _tmp_files[sid] = tmp.name
    node_idx = d2.getNodeIndex(source_node_id)
    source_type_map = {"CONCEN": 1, "MASS": 2, "FLOWPACED": 3, "SETPOINT": 4}
    stype = source_type_map.get(source_type, 1)
    d2.setNodeSourceType(node_idx, stype)
    d2.setNodeSourceQuality(node_idx, concentration)
    res = d2.getComputedTimeSeries()
    times = [int(t) for t in res.Time]
    node_ids = d2.getNodeNameID()
    quality = {node_ids[i]: [round(float(v), 4) for v in res.NodeQuality[:, i]]
               for i in range(len(node_ids))}
    return {"scenario_id": sid, "source_node": source_node_id,
            "concentration": concentration, "times_s": times, "node_quality": quality}


@mcp.tool()
def create_demand_perturbation(network_id: str, node_demands: dict,
                               scenario_id: str = None) -> dict:
    """
    Apply demand multipliers to a set of nodes and simulate.
    node_demands: map of node_id -> multiplier (e.g. {"J1": 2.0}).
    """
    import epyt
    src = _get_session(network_id)
    sid = scenario_id or f"demand_{uuid.uuid4().hex[:6]}"
    tmp = tempfile.NamedTemporaryFile(suffix=".inp", delete=False)
    tmp.close()
    src.saveInputFile(tmp.name)
    d2 = epyt.epanet(tmp.name, verbose=False)
    _sessions[sid] = d2
    _tmp_files[sid] = tmp.name
    for node_id, multiplier in node_demands.items():
        idx = d2.getNodeIndex(node_id)
        base = d2.getNodeBaseDemands()[1][idx - 1]
        d2.setNodeBaseDemands(idx, base * multiplier)
    res = d2.getComputedHydraulicTimeSeries()
    times = [int(t) for t in res.Time]
    node_ids = d2.getNodeNameID()
    pressures = {node_ids[i]: [round(float(v), 3) for v in res.Pressure[:, i]]
                 for i in range(len(node_ids))}
    return {"scenario_id": sid, "times_s": times, "pressures": pressures}


@mcp.tool()
def create_multi_failure_scenario(network_id: str, failed_pipes: list = None,
                                  failed_pumps: list = None, scenario_id: str = None) -> dict:
    """Simulate simultaneous pipe and/or pump failures (closure)."""
    import epyt
    src = _get_session(network_id)
    sid = scenario_id or f"failure_{uuid.uuid4().hex[:6]}"
    tmp = tempfile.NamedTemporaryFile(suffix=".inp", delete=False)
    tmp.close()
    src.saveInputFile(tmp.name)
    d2 = epyt.epanet(tmp.name, verbose=False)
    _sessions[sid] = d2
    _tmp_files[sid] = tmp.name
    for pid in (failed_pipes or []):
        idx = d2.getLinkIndex(pid)
        d2.setLinkInitialStatus(idx, 0)
    for puid in (failed_pumps or []):
        idx = d2.getLinkIndex(puid)
        d2.setLinkInitialStatus(idx, 0)
    res = d2.getComputedHydraulicTimeSeries()
    times = [int(t) for t in res.Time]
    node_ids = d2.getNodeNameID()
    pressures = {node_ids[i]: [round(float(v), 3) for v in res.Pressure[:, i]]
                 for i in range(len(node_ids))}
    return {"scenario_id": sid, "failed_pipes": failed_pipes or [],
            "failed_pumps": failed_pumps or [], "times_s": times, "pressures": pressures}


# ── Info / health routes ───────────────────────────────────────────────────────

async def root(request: Request):
    return JSONResponse({
        "name": "EPANET MCP Server",
        "version": "2.0.0",
        "transports": {
            "sse": "/sse",                  # Claude Desktop (mcp-remote), LM Studio
            "streamable_http": "/mcp",      # HuggingChat, OpenAI Responses API, Cursor
        },
        "tools": 23,
        "loaded_networks": list(_sessions.keys()),
    })


async def health(request: Request):
    return JSONResponse({"status": "ok", "loaded_networks": len(_sessions)})



# ── MCP transport dispatcher ──────────────────────────────────────────────────

class _MCPDispatcher:
    """Routes /mcp → Streamable HTTP transport, all other paths → SSE transport."""

    def __init__(self, sse_app, streamable_app):
        self.sse_app = sse_app
        self.streamable_app = streamable_app

    async def __call__(self, scope, receive, send):
        path = scope.get("path", "")
        if path == "/mcp" or path.startswith("/mcp/"):
            await self.streamable_app(scope, receive, send)
        else:
            await self.sse_app(scope, receive, send)


# ── Streamable HTTP path adapter ──────────────────────────────────────────────
class _StreamableHTTPAdapter:
    """Restores the /mcp path in scope before passing to session_manager.handle_request.

    Starlette strips the mount prefix ('/mcp') when dispatching to child apps,
    but the MCP SDK's StreamableHTTPSessionManager checks scope['path'] == '/mcp'
    internally, so we inject it back.
    """
    def __init__(self, handler):
        self.handler = handler

    async def __call__(self, scope, receive, send):
        if scope.get("type") == "http":
            scope = {**scope, "path": "/mcp", "raw_path": b"/mcp"}
        await self.handler(scope, receive, send)


# ── Build the Starlette app ────────────────────────────────────────────────────

# SSE transport  →  GET /sse  +  POST /messages/
_mcp_sse_app = mcp.sse_app()

# Streamable HTTP transport  →  POST/GET /mcp
# Calling streamable_http_app() initialises the session_manager lazily.
_mcp_streamable_app = mcp.streamable_http_app()
from contextlib import asynccontextmanager

@asynccontextmanager
async def lifespan(app):
    """Start the Streamable HTTP session manager on startup."""
    async with mcp.session_manager.run():
        yield

app = Starlette(
    lifespan=lifespan,
    routes=[
        Route("/", root),
        Route("/health", health),
        # Combined MCP dispatcher: routes /mcp → Streamable HTTP, everything else → SSE.
        # Two separate Mount("/") entries don't work in Starlette (first-match wins),
        # so we use a single ASGI dispatcher that inspects the path itself.
        Mount("/", app=_MCPDispatcher(_mcp_sse_app, _mcp_streamable_app)),
    ]
)

app = APIKeyMiddleware(app)

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

# ── Entrypoint ─────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860, timeout_keep_alive=75)