#!/usr/bin/env python3 """Run one isolated online-model day through MCP and materialize noisy SCADA.""" from __future__ import annotations import argparse import asyncio import json import os import sys from datetime import datetime from pathlib import Path def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--start", default="2026-05-01T00:00:00+08:00") parser.add_argument( "--generation-state", type=Path, default=Path("data/benchmark-generation/tjwater2-20261105/state.json"), ) parser.add_argument( "--state", type=Path, default=Path("data/benchmark-generation/tjwater2-20261105/normal-day-state.json"), ) parser.add_argument( "--design", type=Path, default=Path("resources/benchmarks/tjwater2-v2-burst-balanced.json"), ) parser.add_argument("--agent-env", type=Path, default=Path("../next-tjwater-agent/.local.env")) parser.add_argument("--noise-seed", default="tjwater2-normal-scada-v1") args = parser.parse_args() sys.path.insert(0, str(Path(__file__).resolve().parents[1])) os.umask(0o077) from dotenv import load_dotenv from fastapi.testclient import TestClient from psycopg import connect from psycopg.conninfo import make_conninfo from psycopg.rows import dict_row from app.main import app from scripts.benchmark.calibration import save_json from scripts.benchmark.generation import write_realtime_scada_day from scripts.benchmark.simulation_design import load_simulation_design from scripts.check_burst_simulation_design import resolve_project from scripts.run_burst_benchmark_mcp import exchange_token, execute_batch, mcp_call load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False) start = datetime.fromisoformat(args.start) if start.tzinfo is None or start.minute % 15 or start.second or start.microsecond: raise ValueError("start must be timezone-aware and aligned to 15 minutes") generation = json.loads(args.generation_state.read_text(encoding="utf-8")) if generation.get("status") != "prepared" or not generation.get("project_id"): raise ValueError("isolated benchmark project is not prepared") state_path = args.state.resolve() state_path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) state = json.loads(state_path.read_text(encoding="utf-8")) if state_path.exists() else { "version": "normal-scada-day@1", "target": generation["target"], "project_id": generation["project_id"], "start": start.isoformat(), "duration_seconds": 86400, "noise_seed": args.noise_seed, "status": "preparing", } expected = { "target": generation["target"], "project_id": generation["project_id"], "start": start.isoformat(), "noise_seed": args.noise_seed, } if any(state.get(key) != value for key, value in expected.items()): raise ValueError("normal-day state belongs to different inputs") save_json(state_path, state) design = load_simulation_design(args.design) source_biz, source_iot = asyncio.run(resolve_project(design)) business_dsn = make_conninfo(source_biz, dbname=state["target"]) timescale_dsn = make_conninfo(source_iot, dbname=state["target"]) with connect(business_dsn, row_factory=dict_row) as connection: node = connection.execute( """SELECT node_id FROM asset.scada_devices WHERE device_type='pressure' AND node_id IS NOT NULL ORDER BY device_id LIMIT 1""" ).fetchone()["node_id"] link = connection.execute( """SELECT link_id FROM asset.scada_devices WHERE device_type='pipe_flow' AND link_id IS NOT NULL ORDER BY device_id LIMIT 1""" ).fetchone()["link_id"] if not state.get("job_id"): token = exchange_token(args.agent_env) with TestClient(app, base_url="http://127.0.0.1:8000") as client: headers = { "Authorization": "Bearer " + token, "X-Project-Id": state["project_id"], } binding = client.post( "/api/v1/mcp-query-contexts", headers=headers, json={"native_session_id": "benchmark-normal-scada-v1", "purpose": "ordinary_chat"}, ) binding.raise_for_status() queued = mcp_call(client, token, state["project_id"], "run_simulation", { "context_handle": binding.json()["context_handle"], "timeout_seconds": 240, "start_time": start.isoformat(), "duration": 86400, "output_node_ids": [node], "output_link_ids": [link], }) state.update({"job_id": queued["job_id"], "status": queued["status"]}) save_json(state_path, state) print({"job_id": state["job_id"], "status": state["status"]}, flush=True) async def stored_job(include_result: bool = False): from app.services.calculation_worker import queue_repository async with queue_repository(state["project_id"]) as repository: return await repository.get(state["job_id"], include_result=include_result) stored = asyncio.run(stored_job()) if stored["status"] == "queued": asyncio.run(execute_batch(state["project_id"], [state["job_id"]], 1)) stored = asyncio.run(stored_job()) if stored["status"] != "succeeded": raise RuntimeError(f"normal online simulation ended with {stored['status']}") result = asyncio.run(stored_job(include_result=True))["result"] calculation_data = (result or {}).get("data") or {} scenario_result = calculation_data.get("result") or calculation_data units = scenario_result.get("units") or {} if str(units.get("pressure", "")).upper() not in {"MTR", "M", "METERS"}: raise ValueError(f"unexpected model pressure unit: {units.get('pressure')}") if str(units.get("flow", "")).upper() != "LPS": raise ValueError(f"unexpected model flow unit: {units.get('flow')}") if not state.get("scada"): state["status"] = "simulated" save_json(state_path, state) state["scada"] = write_realtime_scada_day( business_dsn=business_dsn, timescale_dsn=timescale_dsn, start=start, noise_seed=args.noise_seed, ) state["model_units"] = units state["status"] = "scada_materialized" save_json(state_path, state) print(json.dumps({"status": state["status"], "job_id": state["job_id"], "scada": state["scada"]}, ensure_ascii=False), flush=True) return 0 if __name__ == "__main__": raise SystemExit(main())