""" 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)