"""MCP domain dispatch reuses Server repositories and services, never an HTTP loopback.""" import asyncio import math from contextlib import nullcontext from typing import Any from pydantic import BaseModel from app.mcp.scope import QueryScope from app.mcp.catalog import QueryTool, TOOLS from app.mcp.models import BoundContext from datetime import datetime, timezone from uuid import uuid4 from fastapi import HTTPException from fastapi.encoders import jsonable_encoder from app.auth.permissions import permissions_for_context from app.infra.db.postgresql import read_queries as reads from app.infra.db.postgresql.spatial_queries import nearby_elements from app.infra.db.postgresql.analysis import AnalysisRepository from app.infra.db.postgresql.gis_attributes import GisAttributeRepository from app.infra.db.postgresql.network_settings import NetworkSettingsRepository from app.infra.db.postgresql.operational_tasks import OperationalTaskRepository from app.services.operational_tasks import OperationalTaskService from app.mcp.models import QueryResult, PageInfo from app.mcp.scenario_service import SCENARIO_TOOLS from app.mcp.work_order_skills import read_work_order_skill def finite_data(value: Any) -> tuple[Any, int]: """JSON cannot represent infinity/NaN; preserve the gap as null with provenance.""" if isinstance(value, float) and not math.isfinite(value): return None, 1 count = 0 if isinstance(value, dict): result = {} for key, item in value.items(): result[key], replaced = finite_data(item) count += replaced return result, count if isinstance(value, list): result = [] for item in value: clean, replaced = finite_data(item) result.append(clean) count += replaced return result, count return value, 0 def task_service(scope: QueryScope) -> OperationalTaskService: return OperationalTaskService( OperationalTaskRepository(scope.connection, scope.project.project_id), scope.project, set(permissions_for_context(scope.project)), ) def page_result(rows, args): has_more = len(rows) > args.limit return ( rows[: args.limit], PageInfo(limit=args.limit, offset=args.offset, total=None, has_more=has_more), has_more, ) async def authorize_query( scope: QueryScope, binding: BoundContext, tool: QueryTool, args: BaseModel ) -> None: if tool.name not in binding.tool_cap: raise HTTPException(403, "Query is outside the captured tool catalog") permissions = permissions_for_context(scope.project) if tool.mutating: if binding.purpose != "ordinary_chat" or binding.task_id is not None or scope.project.project_role != "member": raise HTTPException(403, "Writes require a project member's ordinary interactive chat") required = {tool.permission} if tool.name == "save_simulation_scenario": required.add("scada.view") if not required.issubset(permissions) or not required.issubset(binding.permission_cap): raise HTTPException(403, "Write and input data permissions are required") if tool.computation: if tool.name == "optimize_sensor_placement" and scope.project.project_role != "member": raise HTTPException(403, "Sensor placement requires project member access") if binding.purpose == "task_reader": raise HTTPException(403, "Read-only task agents cannot start calculations") required = {tool.permission} if tool.name in SCENARIO_TOOLS: required.add("scada.view") if tool.name in {"run_burst_detection", "run_burst_location", "run_leakage_identification"}: required.add("simulation.view" if getattr(args, "data_source", "monitoring") == "simulation" else "scada.view") if not required.issubset(permissions) or not required.issubset(binding.permission_cap): raise HTTPException(403, "Calculation and input data permissions are required") permission = ( "scada.view" if tool.name == "get_network_schema" and args.kind == "scada-device" else tool.permission ) if tool.task_only and binding.task_id is None: raise HTTPException(403, "This query requires a bound task") granted = permission in permissions and permission in binding.permission_cap source = { "list_scada_devices": "scada_devices", "get_scada_device": "scada_devices", "query_scada_readings": "scada_readings", "query_element_history": "scada_readings", "get_workbench_runtime": "workbench_runtime", }.get(tool.name) if tool.name == "query_peripheral_data": source = args.source if not granted and binding.task_id and source in binding.task_scope_cap: async with scope.connection.transaction(): granted = await task_service(scope).can_read_source(binding.task_id, source) if not granted: raise HTTPException(403, "Data source permission is required") if ( tool.name == "query_element_history" and args.mode != "observed" and ( "simulation.view" not in permissions or "simulation.view" not in binding.permission_cap ) ): raise HTTPException(403, "Model comparison permission is required") async def execute_query( scope: QueryScope, binding: BoundContext, tool: QueryTool, args: BaseModel, *, full_snapshot: bool = False, ) -> QueryResult: await authorize_query(scope, binding, tool, args) if tool.mutating: from app.mcp.mutation_service import execute_mutation return await execute_mutation(scope, binding, tool, args) from app.mcp.extended_service import EXTENDED_READ_TOOLS, execute_extended if tool.name in EXTENDED_READ_TOOLS: return await execute_extended(scope, binding, tool, args) name = tool.name conn = scope.connection page = None truncated = False units = {} origin = "system_data" gaps = [] source_ref = "mcp:" + str(uuid4()) evidence_id = None if tool.computation: from app.mcp.analysis_service import TIMESCALE_CALCULATIONS, calculate # Context/receipt reads may have started an implicit transaction. Release # it before potentially long CPU/IO work; evidence uses a new transaction. await conn.commit() async with ( scope.simulation_routing() if name in TIMESCALE_CALCULATIONS else nullcontext() ): data = await calculate( name, args, network=scope.project.project_code, username=scope.user.username, ) data = { "parameters": args.model_dump(mode="json", exclude={"context_handle"}), "result": data, } gaps = [{ "run_simulation": "Online hydraulic recalculation updated the realtime baseline; results cover requested elements only, not observed SCADA data. No standalone analysis run was saved.", "run_contaminant_simulation": "Forward model simulation for an assumed source, not source inversion or observed contamination; only requested elements are returned.", "analyze_valve_closure": "Simulated hydraulic impact for requested elements only; no physical valves were operated.", "analyze_flushing": "Simulated flushing hydraulics only; no field action or water-quality clearance was verified.", "analyze_pressure_regulation": "Evaluation of a supplied pump schedule, not automatic optimization or device control; requested elements only.", "optimize_sensor_placement": "Candidate pressure monitoring locations, not a saved deployment scheme or a guarantee of contaminant source observability.", "run_leakage_identification": "DMA leakage estimates depend on the total leakage constraint and aligned model/SCADA time; they do not confirm an exact leak location.", "run_burst_detection": "Pressure anomalies alone do not confirm a burst or locate a pipe.", "run_burst_location": "Candidate ranking depends on the leakage assumption, model and sensor coverage; field confirmation is required.", "analyze_valve_isolation": "Topology-only isolation advice; valve operability, hydraulic impact and field actions are not verified or executed.", "calculate_demand_allocation": "Pipe-length-weighted allocation only; model demands were not modified. Units follow the current model.", "run_burst_simulation": "Forward hydraulic simulation for assumed orifice areas, not detection or confirmed burst location; only requested elements returned.", }[name]] if name == "calculate_demand_allocation": truncated = data["result"]["has_more"] or args.offset > 0 page = PageInfo(limit=args.limit, offset=args.offset, total=data["result"]["total"], has_more=data["result"]["has_more"]) if getattr(args, "data_source", None) == "simulation": gaps.append("Input is an existing analysis run and model baseline, not observed SCADA data.") elif name == "query_nearby_elements": rows = await nearby_elements( conn, args.element_type, center=args.center, radius_m=args.radius_m, limit=args.limit + 1, offset=args.offset, ) items, page, truncated = page_result(rows, args) truncated = truncated or args.offset > 0 data = { "source": "model_network", "element_type": args.element_type, "center": list(args.center), "center_crs": "EPSG:4326", "radius_m": args.radius_m, "distance_method": "WGS84_spheroid", "items": items, } units = {"distance_m": "m", "radius_m": "m"} gaps = ["Spatial proximity does not establish hydraulic connectivity or a supply relationship"] elif name in {"get_element_properties", "list_elements"}: rows = await reads.asset_rows( conn, args.element_type, element_id=getattr(args, "element_id", None), limit=getattr(args, "limit", 1) + 1, offset=getattr(args, "offset", 0), ) if name == "get_element_properties": if not rows: raise HTTPException(404, "Element was not found") data = rows[0] else: data, page, truncated = page_result(rows, args) elif name in {"list_scada_devices", "get_scada_device"}: rows = await reads.scada_devices( conn, device_id=getattr(args, "device_id", None), limit=getattr(args, "limit", 1) + 1, offset=getattr(args, "offset", 0), ) if name == "get_scada_device": if not rows: raise HTTPException(404, "SCADA device was not found") data = rows[0] else: data, page, truncated = page_result(rows, args) elif name == "get_network_schema": from app.services import tjnetwork mapping = { "option": "get_option_v3_schema", "scada-device": "get_scada_info_schema", "pump-energy": "get_pump_energy_schema", } function = getattr( tjnetwork, mapping.get(args.kind, "get_" + args.kind.replace("-", "_") + "_schema") ) data = await asyncio.to_thread(function, scope.project.project_code) elif name == "get_network_options": if args.kind == "pump-energy": rows = await reads.read_rows( conn, "SELECT price,pattern_id AS pattern,efficiency_curve_id AS effic FROM network.pump_energy_settings WHERE pump_id=%s", (args.pump_id,), ) data = {"pump": args.pump_id, "price": None, "pattern": None, "effic": None} if rows: data.update(rows[0]) else: table = { "time": "time_settings", "energy": "energy_settings", "network": "simulation_settings", }[args.kind] where = " WHERE engine_version='v3'" if args.kind == "network" else "" rows = await reads.read_rows( conn, "SELECT key,value FROM network." + table + where + " ORDER BY key LIMIT 501", ) data = {str(row["key"]): str(row["value"]) for row in rows[:500]} truncated = len(rows) > 500 elif name in { "query_scada_readings", "query_model_baseline_series", "get_model_baseline_at_time", "query_model_baseline_field", "query_analysis_series", "query_analysis_values", }: source = ( "scada" if name == "query_scada_readings" else "analysis" if name.startswith("query_analysis") else "realtime" ) ids = ( getattr(args, "device_ids", None) or getattr(args, "element_ids", None) or ([args.element_id] if hasattr(args, "element_id") else []) ) async with scope.timescale() as ts: rows = await reads.series_rows( ts, source=source, element_type=getattr(args, "element_type", "junction"), ids=ids, start=getattr(args, "start", None), end=getattr(args, "end", None), time=getattr(args, "time", None), field=getattr(args, "field", None), run_id=getattr(args, "run_id", None), limit=getattr(args, "limit", 1) + 1, offset=getattr(args, "offset", 0), ) if name == "get_model_baseline_at_time": data = rows[:1] else: data, page, truncated = page_result(rows, args) if source != "scada": async with conn.transaction(): units = await NetworkSettingsRepository.get_result_units(conn) if source == "analysis": run = await AnalysisRepository.get_run(conn, args.run_id) if not run: raise HTTPException(404, "Analysis run was not found") recorded = (run.get("parameters") or {}).get("result_units") if isinstance(recorded, dict): units = recorded else: gaps.append( "Analysis units inferred from current network settings" ) if hasattr(args, "field") and args.field not in units: gaps.append(f"Unit metadata is unavailable for field {args.field}") else: metadata = [ r for device in ids for r in await reads.scada_devices(conn, device_id=device, limit=1) ] units = { r["device_id"]: r["measurement_unit"] for r in metadata if r["device_id"] in ids } elif name in {"list_analysis_runs", "get_analysis_run", "list_analysis_results"}: async with conn.transaction(): if name == "get_analysis_run": data = await AnalysisRepository.get_run(conn, args.run_id) if not data: raise HTTPException(404, "Analysis run was not found") elif name == "list_analysis_runs": # Fetch one additional row with the repository's supported summary-page bound. rows = await AnalysisRepository.list_run_summaries( conn, limit=args.limit + 1, offset=args.offset ) data, page, truncated = page_result(rows, args) else: rows = await AnalysisRepository.list_results( conn, args.run_id, args.result_type, limit=args.limit, lookahead=True, offset=args.offset, ) data, page, truncated = page_result(rows, args) elif name in {"list_sensor_placement_runs", "get_sensor_placement_run"}: rows = await reads.sensor_runs( conn, run_id=getattr(args, "run_id", None), limit=getattr(args, "limit", 1) + 1, offset=getattr(args, "offset", 0), ) if name == "get_sensor_placement_run": if not rows: raise HTTPException(404, "Sensor placement run was not found") data = rows[0] else: data, page, truncated = page_result(rows, args) elif name == "query_element_history": from app.domain.schemas.timeseries_history import ElementHistoryQuery from app.services.timeseries_history import TimeseriesHistoryService query = ElementHistoryQuery( start_time=args.start, end_time=args.end, mode=args.mode, run_id=args.run_id, elements=[e.model_dump() for e in args.elements], ) async with scope.timescale() as ts, conn.transaction(), ts.transaction(): result = await TimeseriesHistoryService.query( ts, conn, query, row_limit=args.limit, row_offset=args.offset ) data = result.model_dump(mode="json") truncated = result.truncated page = PageInfo(limit=args.limit, offset=args.offset, has_more=truncated) truncated = truncated or args.offset > 0 elif name in {"get_element_gis_properties", "get_source_gis_properties"}: async with conn.transaction(): if name == "get_element_gis_properties": data = await GisAttributeRepository.list_element_groups( conn, args.element_type, args.element_id ) else: data = await GisAttributeRepository.get_source_feature_group( conn, args.layer_key, args.source_fid ) if data is None: raise HTTPException(404, "Source GIS feature was not found") elif name == "get_task_context": if args.section in {"summary", "skill"}: data = await reads.task_context_summary(conn, project_id=binding.project_id, task_id=binding.task_id, report_revision=args.report_revision) skill = read_work_order_skill( data.get("task", {}).get("workflow"), include_content=args.section == "skill" ) if args.section == "skill": data = {"section": "skill", "skill": skill} else: data["handling_skill"] = skill data["display_time_zone"] = "Asia/Shanghai" data["approved_data_scopes"] = binding.task_scope_cap data["bound_run_id"] = str(binding.run_id) if binding.run_id else None data["available_sections"] = ["task", "report", "evidence", "actions", "progress", "access_requests", "plan", "run", "skill"] truncated = True gaps.append("Task context is a bounded summary; use section, data_path, limit and offset for detail. Previews are not complete reports.") else: projection = await reads.task_context_page(conn, project_id=binding.project_id, task_id=binding.task_id, section=args.section, report_revision=args.report_revision, data_path=args.data_path, limit=args.limit, offset=args.offset) total = projection["total"] if args.section in {"evidence", "actions", "access_requests"} and not args.data_path: total = projection["snapshot_metadata"].get("total", total) if projection["value_type"] in {"array", "object", "string"}: page = PageInfo(limit=args.limit, offset=args.offset, total=total, has_more=args.offset + args.limit < total) elif args.offset: raise HTTPException(422, "Scalar task detail cannot use an offset") truncated = bool(projection.get("nested") or args.offset or page and page.has_more) data = {"section": args.section, "selected_path": args.data_path, "value": projection["value"], "metadata": projection["snapshot_metadata"]} if projection.get("nested"): gaps.append("Nested values have data_path descriptors; select a descriptor path for its values") elif name == "get_task_evidence": projection = await reads.evidence_page( conn, project_id=scope.project.project_id, task_id=binding.task_id, evidence_id=args.evidence_id, data_path=args.data_path, limit=args.limit, offset=args.offset, mcp_sources=list(TOOLS), full_snapshot=full_snapshot, ) document = projection if full_snapshot else projection["snapshot_metadata"] origin = document["origin"] source_ref = document["source_ref"] evidence_id = args.evidence_id truncated = document.get("truncated", False) if full_snapshot: data = document else: selected = projection["value"] if projection["value_type"] in {"array", "object", "string"}: total = projection["total"] page = PageInfo( limit=args.limit, offset=args.offset, total=total, has_more=args.offset + args.limit < total, ) truncated = truncated or page.has_more or args.offset > 0 elif args.offset: raise HTTPException(422, "Scalar evidence cannot use an offset") if projection.get("nested"): truncated = True gaps.append( "Nested collections have data_path descriptors; query a descriptor path for their values" ) data = { "snapshot_metadata": document, "query_metadata": projection["query_metadata"], "selected_path": args.data_path, "value": selected, } elif name == "get_workbench_runtime": from app.services.workbench_runtime import read_workbench_runtime data = (await read_workbench_runtime(scope.project, conn)).model_dump( mode="json" ) elif name == "query_peripheral_data": from app.mcp.peripheral import read_peripheral_query snapshot = await read_peripheral_query(scope, args) data = snapshot.data origin = snapshot.origin truncated = snapshot.truncated else: raise ValueError("query handler is not registered") window = ( {"start": args.start, "end": args.end} if hasattr(args, "start") and args.start is not None else None ) data, nonfinite = finite_data(jsonable_encoder(data)) semantics = ( "comparison" if name == "query_element_history" and args.mode != "observed" else tool.semantics ) return QueryResult( source=name, source_ref=source_ref, evidence_id=evidence_id, queried_at=datetime.now(timezone.utc), semantics=semantics, data=data, data_gaps=gaps + ( [f"{nonfinite} non-finite numeric values were returned as null"] if nonfinite else [] ), page=page, window=window, units=units, truncated=truncated, origin=origin, )