"""Previewed writes with project permissions and a durable no-retry journal.""" import asyncio import hashlib import json import math from datetime import datetime, timezone from uuid import uuid4 from typing import Any, TYPE_CHECKING from fastapi import HTTPException from fastapi.encoders import jsonable_encoder from psycopg import sql from pydantic import BaseModel from app.infra.db.postgresql.mcp_queries import McpQueryRepository from app.infra.db.postgresql.read_queries import read_rows, scada_devices from app.mcp.extended_service import CONFIG_FUNCTIONS, configuration from app.mcp.errors import CalculationCapacityExceeded from app.mcp.models import QueryResult, BoundContext from app.mcp.scope import QueryScope if TYPE_CHECKING: from app.mcp.catalog import QueryTool def digest(value: Any) -> str: return hashlib.sha256(json.dumps(jsonable_encoder(value), sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()).hexdigest() def check_digest(args: BaseModel, preview: dict) -> None: if not args.dry_run and args.expected_digest != digest(preview): raise HTTPException(409, "Preview changed; request a new dry_run before execution") def mutation_arguments(args: BaseModel, *, preview: bool) -> dict[str, Any]: excluded = {"context_handle", "timeout_seconds"} if preview: excluded.update({"dry_run", "expected_digest"}) payload = args.model_dump(mode="json", exclude=excluded) if hasattr(args, "records"): # An omitted cleaned_value means keep it during update; explicit null # means clear it. Preserve that difference in previews and replay keys. payload["records"] = [record.model_dump(mode="json", exclude_unset=True) for record in args.records] return payload def check_properties(schema: dict, properties: dict, *, creating: bool = False) -> None: def finite(value): if isinstance(value, float) and not math.isfinite(value): raise HTTPException(422, "Model values must be finite") if isinstance(value, dict): for child in value.values(): finite(child) elif isinstance(value, list): for child in value: finite(child) finite(properties) for key, value in properties.items(): spec = schema.get(key) if spec is None or spec.get("readonly"): raise HTTPException(422, "Unknown or read-only model property: " + key) if value is None and not spec.get("optional", False): raise HTTPException(422, "Required model property cannot be null: " + key) expected = spec.get("type") if value is not None: valid = {"str": isinstance(value, str), "float": isinstance(value, (int, float)) and not isinstance(value, bool), "int": isinstance(value, int) and not isinstance(value, bool), "bool": isinstance(value, bool), "list": isinstance(value, list), "str_list": isinstance(value, list) and all(isinstance(v, str) for v in value), "float_list": isinstance(value, list) and all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in value), "tuple_list": isinstance(value, list)}.get(expected, True) if not valid: raise HTTPException(422, "Invalid model property type: " + key) if isinstance(value, list) and len(value) > 500: raise HTTPException(422, "Model property exceeds 500 items") if isinstance(value, list) and isinstance(spec.get("element"), dict): for item in value: if not isinstance(item, dict): raise HTTPException(422, "Model list item must be an object") check_properties(spec["element"], item, creating=True) if creating: required = {key for key, spec in schema.items() if not spec.get("readonly") and not spec.get("optional")} if not required.issubset(properties): raise HTTPException(422, "Missing required model properties: " + ",".join(sorted(required - properties.keys()))) def model_change(network: str, name: str, args: BaseModel, *, execute: bool) -> dict[str, Any]: from app.services import tjnetwork from app.native.wndb.core.database import ChangeSet, model_mutation_transaction # Preview revalidation and the actual write share the same model mutation # lock; another WNDB writer cannot change the target between them. with model_mutation_transaction(network): if name == "edit_model_element": kind, identity, object_id = args.element_type, "id", args.element_id current = getattr(tjnetwork, "get_" + kind)(network, object_id) schema = getattr(tjnetwork, "get_" + kind + "_schema")(network) else: kind, identity = CONFIG_FUNCTIONS[args.kind] object_id = args.object_id current = configuration(network, args.kind, object_id) schema = getattr(tjnetwork, "get_" + kind + "_schema")(network) check_properties(schema, args.properties, creating=args.operation == "create") if identity and args.operation == "create" and current: raise HTTPException(409, "Object already exists") if identity and args.operation != "create" and not current: raise HTTPException(404, "Model object not found") preview = {"tool": name, "operation": args.operation, "kind": kind, "object_id": object_id, "before": current, "properties": args.properties} if not execute: return preview check_digest(args, preview) prefix = {"create": "add_", "update": "set_", "delete": "delete_"}[args.operation] values = ({identity: object_id} if identity else {}) | args.properties if kind == "region" and "boundary" in values: values["boundary"] = [tuple(point) for point in values["boundary"]] result = getattr(tjnetwork, prefix + kind)(network, ChangeSet(values)) if not result.operations: raise HTTPException(409, "Model rejected the change; no operations were applied") return {"operations": result.operations, "field_device_control": False} async def scada_snapshot(scope: QueryScope, args: BaseModel, *, connection: Any = None, lock: bool = False) -> dict[str, Any]: ids = sorted({r.device_id for r in args.records}) if hasattr(args, "records") else sorted(set(args.device_ids)) for device_id in ids: if not await scada_devices(scope.connection, device_id=device_id, limit=1): raise HTTPException(404, "SCADA device not found in the current project") async def read(ts): suffix = " FOR UPDATE" if lock else "" if hasattr(args, "records"): rows = [] for record in args.records: rows.extend(await read_rows(ts, "SELECT time,device_id,monitored_value,cleaned_value FROM scada.measurements WHERE device_id=%s AND time=%s" + suffix, (record.device_id, record.time))) else: rows = await read_rows(ts, """SELECT time,device_id,monitored_value,cleaned_value FROM scada.measurements WHERE device_id=ANY(%s) AND time >= %s AND time <= %s ORDER BY device_id,time LIMIT 10001""" + suffix, (ids, args.start, args.end)) if len(rows) > 10000: raise HTTPException(413, "Narrow the SCADA window to at most 10000 records") return {"count": len(rows), "devices": ids, "state_digest": digest(rows)} if connection is not None: return await read(connection) async with scope.timescale() as ts: return await read(ts) async def preview_change(scope: QueryScope, binding: BoundContext, name: str, args: BaseModel) -> dict[str, Any]: network = scope.project.project_code if name in {"edit_model_element", "edit_model_configuration"}: from app.mcp.analysis_service import bounded_calculation return await bounded_calculation(model_change, network, name, args, execute=False) payload = mutation_arguments(args, preview=True) state = None if name in {"clean_scada_readings", "write_scada_readings", "delete_scada_readings"}: state = await scada_snapshot(scope, args) elif name in {"update_sensor_placement"}: from app.services.sensor_placement import get_sensor_placement_run, can_edit_sensor_placement state = await asyncio.to_thread(get_sensor_placement_run, network, args.run_id) if not can_edit_sensor_placement(scope.user, state): raise HTTPException(403, "Cannot edit another user's sensor placement") elif name in {"retry_operational_task", "cancel_task_run"}: from app.mcp.query_service import task_service state = (await task_service(scope).detail(args.target_task_id)).model_dump(mode="json") elif name == "update_task_schedule": from app.mcp.query_service import task_service schedules = await task_service(scope).list_schedules() state = next((s.model_dump(mode="json") for s in schedules if s.schedule_id == args.schedule_id), None) if state is None: raise HTTPException(404, "Schedule not found in the current project") return {"tool": name, "parameters": payload, "before": state} def saved_scenario(network: str, username: str, args: BaseModel) -> Any: from app.services import simulation_scenarios as scenarios from app.services.simulation import query_corresponding_element_id_and_query_id parameters = args.parameters.model_dump(exclude_none=True) parameters.pop("output_node_ids") parameters.pop("output_link_ids") common = dict(name=network, modify_pattern_start_time=parameters.pop("start_time").isoformat(), modify_total_duration=parameters.pop("duration"), scheme_name=args.name, persist_results=True) if args.scenario == "burst": return scenarios.burst_analysis(**common, username=username, burst_ID=parameters["burst_pipe_ids"], burst_size=parameters["orifice_areas_cm2"]) if args.scenario == "contaminant": return scenarios.contaminant_simulation(**common, username=username, source=parameters["source"], concentration=parameters["concentration"], source_pattern=parameters.get("pattern")) if args.scenario == "valve_closure": return scenarios.valve_close_analysis(**common, username=username, modify_valve_opening={valve: 0.0 for valve in parameters["valves"]}) if args.scenario == "flushing": from app.native.wndb.model.options import get_option if get_option(network).get("UNITS") not in {"LPS", "CMH"}: raise ValueError("Flushing supports LPS or CMH models only") return scenarios.flushing_analysis(**common, username=username, drainage_node_ID=parameters["drainage_node_id"], flushing_flow=parameters["flushing_flow_m3h"], modify_valve_opening={valve: 0.0 for valve in parameters["closed_valves"]}) mappings = query_corresponding_element_id_and_query_id(network) controls = parameters["pump_control"] if set(controls) - (set(mappings.fixed_pumps) | set(mappings.variable_pumps)): raise ValueError("Unknown model pump") if any(value not in (0, 1) for key, values in controls.items() if key in mappings.fixed_pumps for value in values): raise ValueError("Fixed pumps accept only 0 or 1") return scenarios.pressure_regulation(**common, username=username, modify_fixed_pump_pattern={key: value for key, value in controls.items() if key in mappings.fixed_pumps}, modify_variable_pump_pattern={key: value for key, value in controls.items() if key in mappings.variable_pumps}, modify_tank_initial_level=parameters["tank_init_level"], scada_mappings=mappings) async def apply_change(scope: QueryScope, binding: BoundContext, name: str, args: BaseModel) -> Any: network = scope.project.project_code if name in {"edit_model_element", "edit_model_configuration"}: from app.mcp.analysis_service import bounded_calculation return await bounded_calculation(model_change, network, name, args, execute=True) if name in {"clean_scada_readings", "write_scada_readings", "delete_scada_readings"}: from app.infra.db.timescaledb.repositories.scada import ScadaRepository from app.services.timeseries_analysis import TimeseriesAnalysisService async with scope.timescale() as ts, ts.transaction(): await ts.execute("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE") # Serialize with other MCP maintenance calls, then validate again. await ts.execute("SELECT pg_advisory_xact_lock(hashtextextended(%s,0))", ("mcp-scada-maintenance",)) before = await scada_snapshot(scope, args, connection=ts, lock=True) payload = mutation_arguments(args, preview=True) check_digest(args, {"tool": name, "parameters": payload, "before": before}) if name == "clean_scada_readings": result = await TimeseriesAnalysisService.clean_scada_data(ts, scope.connection, args.device_ids, args.start, args.end) return {"status": result, "field": "cleaned_value"} if name == "delete_scada_readings": async with ts.cursor() as cur: await cur.execute("DELETE FROM scada.measurements WHERE device_id=ANY(%s) AND time >= %s AND time <= %s", (args.device_ids, args.start, args.end)) return {"deleted_rows": cur.rowcount} if args.operation == "insert": await ScadaRepository.insert_scada_batch(ts, [r.model_dump() for r in args.records]) else: for record in args.records: fields = sorted(record.model_fields_set & {"monitored_value", "cleaned_value"}) statement = sql.SQL("UPDATE scada.measurements SET {} WHERE device_id=%s AND time=%s").format( sql.SQL(",").join(sql.SQL("{}=%s").format(sql.Identifier(field)) for field in fields)) async with ts.cursor() as cur: await cur.execute(statement, tuple(getattr(record, field) for field in fields) + (record.device_id, record.time)) if cur.rowcount != 1: raise HTTPException(409, "SCADA update requires exactly one existing record") return {"operation": args.operation, "records": len(args.records), "provenance": args.provenance} if name == "save_simulation_scenario": from app.mcp.analysis_service import bounded_calculation async with scope.simulation_routing(): await bounded_calculation(saved_scenario, network, scope.user.username, args) return {"persisted": True, "name": args.name, "scenario": args.scenario, "results_query": "list_analysis_runs"} if name == "create_sensor_placement": from app.mcp.analysis_service import bounded_calculation from app.services.sensor_placement import optimize_sensor_placement_by_kmeans, optimize_sensor_placement_by_sensitivity optimizer = optimize_sensor_placement_by_sensitivity if args.method == "sensitivity" else optimize_sensor_placement_by_kmeans return await bounded_calculation(optimizer, project_code=network, run_name=args.name, created_by=scope.user.username, sensor_count=args.sensor_count, min_diameter=args.min_diameter, sensor_type="pressure", persist_results=True) if name == "update_sensor_placement": from app.services.sensor_placement import update_sensor_placement_run return await asyncio.to_thread(update_sensor_placement_run, network, args.run_id, **args.changes.model_dump()) from app.mcp.query_service import task_service service = task_service(scope) if name == "create_operational_task": if args.command.agent_session_id is not None or args.command.task_type != "event_investigation" or args.command.source.kind not in {"operator", "agent"}: raise HTTPException(422, "Only unclaimed operator/agent event investigations may be created") return await service.create(args.command) if name == "create_task_schedule": return await service.create_schedule(args.command) if name == "update_task_schedule": return await service.update_schedule(args.schedule_id, args.command) if name == "retry_operational_task": return await service.retry_task(args.target_task_id, args.command) if name == "cancel_task_run": return await service.cancel_run(args.target_task_id, args.run_id, args.command) raise ValueError("Unknown mutation") def retryable_capacity_rejection(result: dict[str, Any]) -> bool: return (result.get("status") == "rejected" and result.get("retryable") is True and result.get("error", {}).get("code") == CalculationCapacityExceeded.code) async def execute_mutation(scope: QueryScope, binding: BoundContext, tool: "QueryTool", args: BaseModel) -> QueryResult: name = tool.name repo = McpQueryRepository(scope.connection, scope.project.project_id) key, call_id = "mutation:" + binding.subject + ":" + binding.native_session_id, str(args.request_id) parameters = {"tool": name, "arguments": mutation_arguments(args, preview=False)} if not args.dry_run: async with scope.connection.transaction(): await repo.lock_receipt(key, call_id) previous = await repo.receipt(key, call_id) if previous: if previous["parameters"] != parameters: raise HTTPException(409, "request_id already has different parameters") if previous["result"].get("status") == "completed": return QueryResult.model_validate(previous["result"]["output"]) if not retryable_capacity_rejection(previous["result"]): raise HTTPException(409, "Operation already started; verify its state before any retry") preview = await preview_change(scope, binding, name, args) if args.dry_run: result_data = {"dry_run": True, "preview": preview, "expected_digest": digest(preview), "request_id": call_id} if len(json.dumps(jsonable_encoder(result_data)).encode()) > 10000: result_data["preview"] = {"tool": name, "summary": "Preview payload omitted; inspect the target with paged read tools", "before_digest": digest(preview.get("before"))} else: check_digest(args, preview) async with scope.connection.transaction(): await repo.lock_receipt(key, call_id) previous = await repo.receipt(key, call_id) if previous: if previous["parameters"] != parameters: raise HTTPException(409, "request_id already has different parameters") if previous["result"].get("status") == "completed": return QueryResult.model_validate(previous["result"]["output"]) if not retryable_capacity_rejection(previous["result"]): raise HTTPException(409, "Operation was started by another request; verify its state") await repo.finish_receipt(key, call_id, {"status": "started"}) else: await repo.save_receipt(key, call_id, parameters, {"status": "started"}) await scope.connection.commit() # Keep started on cancellation/timeout: a native thread may still commit. try: applied = await apply_change(scope, binding, name, args) except CalculationCapacityExceeded as exc: # This specific admission failure proves that no calculation began. # Other errors and cancellation remain ambiguous and are not retried. async with scope.connection.transaction(): await repo.lock_receipt(key, call_id) await repo.finish_receipt(key, call_id, { "status": "rejected", "retryable": True, "error": {"code": exc.code, "message": exc.detail}, "rejected_at": datetime.now(timezone.utc).isoformat(), }) await scope.connection.commit() raise result_data = {"dry_run": False, "request_id": call_id, "result": jsonable_encoder(applied)} result = QueryResult(source=name, source_ref="mcp:" + uuid4().hex, queried_at=datetime.now(timezone.utc), semantics=tool.semantics, data=result_data) if not args.dry_run: async with scope.connection.transaction(): await repo.finish_receipt(key, call_id, {"status": "completed", "output": result.model_dump(mode="json")}) return result