340 lines
20 KiB
Python
340 lines
20 KiB
Python
"""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
|