#!/usr/bin/env python3 """Submit, execute, and materialize isolated burst benchmark cases through MCP.""" from __future__ import annotations import argparse import asyncio import json import os from pathlib import Path import subprocess import sys from uuid import uuid4 def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--design", type=Path, default=Path("resources/benchmarks/tjwater2-v2-burst-balanced.json")) parser.add_argument("--state-dir", type=Path, default=Path("data/benchmark-generation/tjwater2-20261105")) parser.add_argument("--agent-env", type=Path, default=Path("../next-tjwater-agent/.local.env")) parser.add_argument("--limit", type=int, default=1, help="Maximum new cases to complete in this invocation") parser.add_argument("--batches", type=int, default=1, help="Run isolated child processes for this many batches") parser.add_argument("--local-executor-concurrency", type=int, default=1, choices=(1, 2)) args = parser.parse_args() if args.limit < 1 or args.limit > args.local_executor_concurrency: parser.error("--limit must be between 1 and local executor concurrency") if args.batches < 1: parser.error("--batches must be positive") if args.batches > 1: command = [ sys.executable, str(Path(__file__).resolve()), "--design", str(args.design), "--state-dir", str(args.state_dir), "--agent-env", str(args.agent_env), "--limit", str(args.limit), "--local-executor-concurrency", str(args.local_executor_concurrency), "--batches", "1", ] for _ in range(args.batches): subprocess.run(command, check=True) return 0 sys.path.insert(0, str(Path(__file__).resolve().parents[1])) os.umask(0o077) from dotenv import load_dotenv from psycopg.conninfo import make_conninfo from scripts.benchmark.calibration import save_json from scripts.benchmark.generation import write_analysis_scada from scripts.benchmark.simulation_design import load_simulation_design from scripts.check_burst_simulation_design import resolve_project load_dotenv(Path(__file__).resolve().parents[1] / ".env", override=False) state_path = args.state_dir.resolve() / "state.json" state = json.loads(state_path.read_text(encoding="utf-8")) project_id = state.get("project_id") if not project_id or state.get("status") != "prepared": raise ValueError("run isolated preparation before MCP generation") design = load_simulation_design(args.design) pipes = {pipe.pipe_id: pipe for pipe in design.pipes} 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"]) pending = [ row for row in state["schedule"] if state["cases"].get(row["case_key"], {}).get("status") != "scada_materialized" ][:args.limit] for offset in range(0, len(pending), args.local_executor_concurrency): batch = pending[offset:offset + args.local_executor_concurrency] job_ids = submit_batch( batch=batch, pipes=pipes, project_id=project_id, agent_env=args.agent_env, state=state, state_path=state_path, ) asyncio.run(execute_batch(project_id, job_ids, args.local_executor_concurrency)) for row in batch: case = state["cases"][row["case_key"]] run = find_analysis_run(business_dsn, case["scenario_name"]) summary = write_analysis_scada( business_dsn=business_dsn, timescale_dsn=timescale_dsn, run_id=str(run["run_id"]), ) case.update({ "analysis_run_id": str(run["run_id"]), "status": "scada_materialized", "scada": summary, }) save_json(state_path, state) print({"case_key": row["case_key"], "job_id": case["job_id"], "scada_rows": summary["rows"]}, flush=True) return 0 def exchange_token(agent_env: Path) -> str: import httpx from dotenv import dotenv_values from app.mcp.config import McpSettings from scripts.run_operational_worker import CredentialProvider, WorkerSettings values = dotenv_values(agent_env) mcp = McpSettings() async def exchange() -> str: issuer = (values.get("TJWATER_MCP_DEV_HTTP_ISSUER_URI") or values["TJWATER_MCP_ISSUER_URI"]).rstrip("/") async with httpx.AsyncClient(timeout=15) as client: product = await CredentialProvider(WorkerSettings.from_environment()).headers(client) response = await client.post(issuer + "/protocol/openid-connect/token", data={ "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", "client_id": values["TJWATER_MCP_CLIENT_ID"], "client_secret": values["TJWATER_MCP_CLIENT_SECRET"], "subject_token": product["Authorization"].removeprefix("Bearer "), "subject_token_type": "urn:ietf:params:oauth:token-type:access_token", "requested_token_type": "urn:ietf:params:oauth:token-type:access_token", "audience": mcp.audience, "resource": mcp.resource_uri, "scope": "mcp.catalog.read mcp.query mcp.context.bind", }) response.raise_for_status() return response.json()["access_token"] return asyncio.run(exchange()) def mcp_call(client, token: str, project_id: str, name: str, arguments: dict) -> dict: from mcp import types method = "tools/call" call_id = str(uuid4()) meta = { types.PROTOCOL_VERSION_META_KEY: "2026-07-28", types.CLIENT_CAPABILITIES_META_KEY: {}, types.CLIENT_INFO_META_KEY: {"name": "benchmark-generator", "version": "1"}, "io.tjwater/callId": call_id, } response = client.post("/mcp", headers={ "Authorization": "Bearer " + token, "X-Project-Id": project_id, "MCP-Protocol-Version": "2026-07-28", "MCP-Method": method, "MCP-Name": name, "Accept": "application/json, text/event-stream", }, json={ "jsonrpc": "2.0", "id": call_id, "method": method, "params": {"_meta": meta, "name": name, "arguments": arguments}, }) response.raise_for_status() payload = response.json() result = payload.get("result") or {} if payload.get("error") or result.get("isError"): raise RuntimeError("MCP tool call failed") return result["structuredContent"]["data"] def submit_batch(*, batch, pipes, project_id, agent_env, state, state_path) -> list[str]: from fastapi.testclient import TestClient from app.main import app from scripts.benchmark.calibration import save_json token = exchange_token(agent_env) session_id = "benchmark-tjwater2-burst-v1" jobs = [] with TestClient(app, base_url="http://127.0.0.1:8000") as client: common = {"Authorization": "Bearer " + token, "X-Project-Id": project_id} for row in batch: case = state["cases"].setdefault(row["case_key"], {}) request_id = case.setdefault("request_id", str(uuid4())) scenario_name = case.setdefault( "scenario_name", f"bench-{row['case_key']}-{row['event_at'][0:16].replace('-', '').replace(':', '').replace('T', '')}", ) binding = client.post( "/api/v1/mcp-query-contexts", headers=common, json={"native_session_id": session_id, "purpose": "ordinary_chat"}, ) binding.raise_for_status() context_handle = binding.json()["context_handle"] pipe = pipes[row["pipe_id"]] parameters = { "start_time": row["event_at"], "duration": 18000, "output_node_ids": pipe.output_node_ids, "output_link_ids": pipe.output_link_ids, "burst_pipe_ids": [row["pipe_id"]], "orifice_areas_cm2": [row["area_cm2"]], } preview = mcp_call(client, token, project_id, "save_simulation_scenario", { "context_handle": context_handle, "timeout_seconds": 240, "request_id": request_id, "dry_run": True, "name": scenario_name, "scenario": "burst", "parameters": parameters, }) case["expected_digest"] = preview["expected_digest"] save_json(state_path, state) queued = mcp_call(client, token, project_id, "save_simulation_scenario", { "context_handle": context_handle, "timeout_seconds": 240, "request_id": request_id, "dry_run": False, "expected_digest": preview["expected_digest"], "name": scenario_name, "scenario": "burst", "parameters": parameters, }) case.update({"job_id": queued["job_id"], "status": queued["status"]}) jobs.append(queued["job_id"]) save_json(state_path, state) return jobs async def execute_batch(project_id: str, expected_job_ids: list[str], concurrency: int) -> None: from app.services.calculation_worker import CalculationWorker, WorkerSettings, queue_repository worker = CalculationWorker(WorkerSettings( max_concurrency=concurrency, poll_seconds=1, lease_seconds=30, heartbeat_seconds=5, )) rows = [] for expected in expected_job_ids: async with queue_repository(project_id) as repository: await repository.recover() row = await repository.claim(worker.owner, worker.settings.lease_seconds) if row is None or str(row["job_id"]) != expected: raise RuntimeError("isolated benchmark queue order changed unexpectedly") rows.append(row) await asyncio.gather(*(worker.run_job(row) for row in rows)) for row in rows: async with queue_repository(project_id) as repository: stored = await repository.get(row["job_id"]) if stored["status"] != "succeeded": raise RuntimeError(f"benchmark calculation failed with {stored['status']}") def find_analysis_run(business_dsn: str, scenario_name: str) -> dict: from psycopg import connect from psycopg.rows import dict_row with connect(business_dsn, row_factory=dict_row) as connection: rows = connection.execute( """SELECT run_id,name,status,started_at FROM analysis.runs WHERE name=%s ORDER BY created_at DESC LIMIT 2""", (scenario_name,), ).fetchall() if len(rows) != 1 or rows[0]["status"] != "completed": raise ValueError("persisted scenario run is missing, duplicated, or incomplete") return dict(rows[0]) if __name__ == "__main__": raise SystemExit(main())