Files
jiang be6317eb48 feat(server): integrate MCP queries and investigation messaging
Add project-bound OAuth MCP queries and controlled analysis with scoped evidence,
limits, timeouts and a nearby-model spatial query using EPSG:4326 inputs.
Default nearby searches to 200 meters and five results.

Integrate investigation messaging/presentation, health/readiness and typed history;
remove duplicate web search and fix online simulation and compressed baseline writes.
Include project migrations 023-025; keep local credentials and demo seeding excluded.

Validation: 319 focused pytest checks; real PostGIS temporary-table regression;
read-only Jinqian bridge radius comparison; cross-language MCP discovery/calls.
2026-10-19 18:46:58 +08:00

480 lines
23 KiB
Python

import asyncio
from contextlib import asynccontextmanager, contextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock
from uuid import uuid4
import pytest
from fastapi import HTTPException
from pydantic import ValidationError
from app.domain.schemas.operational_data import TaskDataSnapshot
from app.mcp import analysis_service, query_service
from app.mcp.catalog import TOOLS
from app.mcp.scope import QueryScope
from app.infra.db.project_routing import (
ActiveProjectRouting,
activate_project_routing,
get_active_project_routing,
get_project_pgconn_string,
get_project_timescale_pgconn_string,
)
HANDLE = "x" * 43
START = "2026-10-08T12:00:00+08:00"
END = "2026-10-08T13:00:00+08:00"
INPUTS = {
"run_simulation": {"start_time": START, "duration": 3600, "output_node_ids": ["J1"], "output_link_ids": ["L1"]},
"run_contaminant_simulation": {"start_time": START, "output_node_ids": ["J1"], "source": "J2", "concentration": 1.0},
"analyze_valve_closure": {"start_time": START, "output_node_ids": ["J1"], "valves": ["V1"]},
"analyze_flushing": {"start_time": START, "output_node_ids": ["J1"], "drainage_node_id": "J2", "flushing_flow_m3h": 10},
"analyze_pressure_regulation": {"start_time": START, "output_node_ids": ["J1"], "pump_control": {"P1": [1]}},
"optimize_sensor_placement": {"method": "kmeans", "sensor_count": 3},
"run_burst_detection": {"target_time": START, "sensor_nodes": ["1", "2", "3", "4", "5"]},
"run_burst_location": {"data_source": "simulation", "simulation_run_id": str(uuid4()), "burst_leakage": 0.02, "pressure_scada_ids": ["1"], "scada_burst_start": START, "scada_burst_end": END},
"analyze_valve_isolation": {"accident_element": ["540672"]},
"run_leakage_identification": {"scada_start": START, "scada_end": END, "sensor_nodes": ["1"], "q_sum": 0.02, "start_time": 12, "dma_count": 1},
}
def arguments(name, **extra):
return TOOLS[name].model.model_validate({"context_handle": HANDLE, **INPUTS[name], **extra})
@pytest.fixture
def routed_calculation_scope(monkeypatch):
events = []
project = SimpleNamespace(project_id=uuid4(), project_code="logical-project")
business = "host=localhost dbname=physical_business"
timescale = "host=localhost dbname=physical_iot"
async def resolve(project_id, role):
assert project_id == project.project_id
events.append(role)
return SimpleNamespace(
db_type="postgresql" if role == "biz_data" else "timescaledb",
dsn=business if role == "biz_data" else timescale,
)
async def commit():
events.append("metadata_commit")
async def release_business_read():
events.append("business_commit")
metadata = SimpleNamespace(
get_project_db_routing=AsyncMock(side_effect=resolve),
session=SimpleNamespace(commit=AsyncMock(side_effect=commit)),
)
routing = ActiveProjectRouting(project.project_code, business, project_id=project.project_id)
scope = QueryScope(
project, SimpleNamespace(username="actor"),
SimpleNamespace(commit=release_business_read), routing, metadata,
)
monkeypatch.setattr(query_service, "authorize_query", AsyncMock())
return scope, timescale, events
@pytest.mark.parametrize("name", [
"run_simulation", "run_contaminant_simulation", "analyze_valve_closure",
"analyze_flushing", "analyze_pressure_regulation", "run_burst_detection",
"run_burst_location", "run_leakage_identification",
])
def test_timescale_calculation_routes_are_ready_before_native_worker(monkeypatch, routed_calculation_scope, name):
scope, timescale, events = routed_calculation_scope
def worker(tool_name, parameters, network, username):
assert tool_name == name
assert network == scope.project.project_code
assert username == "actor"
assert "context_handle" not in parameters
assert "timeout_seconds" not in parameters
assert get_project_timescale_pgconn_string(network) == timescale
assert get_project_pgconn_string(network) == scope.routing.business_dsn
assert events == ["business_commit", "biz_data", "iot_data", "metadata_commit"]
events.append("worker")
return {"routed": True}
monkeypatch.setattr(analysis_service, "_calculate", worker)
async def execute():
result = await query_service.execute_query(scope, object(), TOOLS[name], arguments(name))
assert get_active_project_routing() is scope.routing
return result
with activate_project_routing(scope.routing):
result = asyncio.run(execute())
assert result.data["result"] == {"routed": True}
assert events[-1] == "worker"
def test_missing_timescale_route_stops_online_simulation_before_model_changes(monkeypatch, routed_calculation_scope):
scope, _, _ = routed_calculation_scope
resolve = scope.metadata.get_project_db_routing.side_effect
async def missing_iot(project_id, role):
return None if role == "iot_data" else await resolve(project_id, role)
scope.metadata.get_project_db_routing.side_effect = missing_iot
calculate = AsyncMock(return_value={})
monkeypatch.setattr(analysis_service, "calculate", calculate)
with activate_project_routing(scope.routing):
with pytest.raises(HTTPException) as denied:
asyncio.run(query_service.execute_query(scope, object(), TOOLS["run_simulation"], arguments("run_simulation")))
assert get_active_project_routing() is scope.routing
assert denied.value.status_code == 503
assert denied.value.detail == "Project TimescaleDB not configured"
calculate.assert_not_awaited()
@pytest.mark.parametrize("name", ["analyze_valve_isolation", "optimize_sensor_placement"])
def test_business_only_calculations_do_not_require_timescale(monkeypatch, routed_calculation_scope, name):
scope, _, events = routed_calculation_scope
def worker(*_):
assert get_active_project_routing() is scope.routing
assert get_project_pgconn_string(scope.project.project_code) == scope.routing.business_dsn
return {"business_only": True}
monkeypatch.setattr(analysis_service, "_calculate", worker)
with activate_project_routing(scope.routing):
result = asyncio.run(query_service.execute_query(scope, object(), TOOLS[name], arguments(name)))
assert result.data["result"] == {"business_only": True}
assert events == ["business_commit"]
scope.metadata.get_project_db_routing.assert_not_awaited()
def test_calculation_failure_restores_business_routing(monkeypatch, routed_calculation_scope):
scope, timescale, _ = routed_calculation_scope
def worker(*_):
assert get_project_timescale_pgconn_string(scope.project.project_code) == timescale
raise RuntimeError("native calculation failed")
monkeypatch.setattr(analysis_service, "_calculate", worker)
async def execute():
with pytest.raises(RuntimeError, match="native calculation failed"):
await query_service.execute_query(scope, object(), TOOLS["run_simulation"], arguments("run_simulation"))
assert get_active_project_routing() is scope.routing
with activate_project_routing(scope.routing):
asyncio.run(execute())
def test_denied_calculation_never_loads_timescale_or_starts_worker(monkeypatch, routed_calculation_scope):
scope, _, events = routed_calculation_scope
monkeypatch.setattr(query_service, "authorize_query", AsyncMock(side_effect=HTTPException(403, "denied")))
calculate = AsyncMock(return_value={})
monkeypatch.setattr(analysis_service, "calculate", calculate)
with pytest.raises(HTTPException) as denied:
asyncio.run(query_service.execute_query(scope, object(), TOOLS["run_simulation"], arguments("run_simulation")))
assert denied.value.status_code == 403
assert events == []
scope.metadata.get_project_db_routing.assert_not_awaited()
calculate.assert_not_awaited()
@pytest.mark.parametrize("name", INPUTS)
def test_catalog_and_snapshots_accept_calculations_but_reject_model_owned_authority(name):
tool = TOOLS[name]
assert tool.computation
assert arguments(name).timeout_seconds == 120
assert tool.protocol_tool().input_schema["additionalProperties"] is False
for field in ("network", "username", "scheme_name", "final_candidates_csv_path", "observed_pressure_data", "persist_results", "simulation_type"):
with pytest.raises(ValidationError):
arguments(name, **{field: "untrusted"})
snapshot = TaskDataSnapshot(source=name, source_ref="mcp:test", queried_at=datetime.now(timezone.utc), origin="system_data", data={}, truncated=False)
assert snapshot.source == name
@pytest.mark.parametrize("name", INPUTS)
def test_calculation_permission_and_reader_boundary(monkeypatch, name):
permissions = {"burst.run", "simulation.run", "simulation.view", "scada.view", "optimization.run"}
monkeypatch.setattr(query_service, "permissions_for_context", lambda _: permissions)
scope = SimpleNamespace(project=SimpleNamespace(project_role="member"))
binding = SimpleNamespace(tool_cap=list(TOOLS), purpose="task_execution", permission_cap=list(permissions), task_id=uuid4(), task_scope_cap=["scada_readings"])
tool, args = TOOLS[name], arguments(name)
asyncio.run(query_service.authorize_query(scope, binding, tool, args))
binding.purpose = "task_reader"
with pytest.raises(HTTPException) as denied:
asyncio.run(query_service.authorize_query(scope, binding, tool, args))
assert denied.value.status_code == 403
binding.purpose = "task_execution"
binding.permission_cap.remove(tool.permission)
with pytest.raises(HTTPException):
asyncio.run(query_service.authorize_query(scope, binding, tool, args))
binding.permission_cap = list(permissions)
permissions.remove(tool.permission)
with pytest.raises(HTTPException):
asyncio.run(query_service.authorize_query(scope, binding, tool, args))
def test_time_sources_and_finite_inputs_are_explicit():
invalid = [
("run_simulation", {"start_time": "2026-10-08T12:00:00"}),
("run_simulation", {"duration": 0}),
("run_simulation", {"duration": 86401}),
("run_simulation", {"timeout_seconds": 0}),
("run_simulation", {"timeout_seconds": 86401}),
("run_simulation", {"timeout_seconds": 30.5}),
("run_simulation", {"timeout_seconds": True}),
("run_simulation", {"output_node_ids": []}),
("run_simulation", {"output_link_ids": [str(i) for i in range(21)]}),
("run_burst_detection", {"target_time": "2026-10-08T12:00:00"}),
("run_burst_detection", {"sensor_nodes": ["1"] * 5}),
("run_burst_location", {"burst_leakage": float("nan")}),
("run_burst_location", {"simulation_run_id": None}),
("run_burst_location", {"data_source": "monitoring"}),
("run_burst_location", {"scada_burst_end": "2026-10-10T12:00:00+08:00"}),
("run_leakage_identification", {"q_sum": 0}),
]
for name, changes in invalid:
with pytest.raises(ValidationError):
arguments(name, **changes)
args = arguments("run_burst_location", data_source="monitoring", simulation_run_id=None,
scada_normal_start="2026-10-07T12:00:00+08:00", scada_normal_end="2026-10-07T13:00:00+08:00")
assert args.data_source == "monitoring"
@pytest.mark.parametrize("overrides", [{}, {
"pressure_scada_ids": None, "scada_normal_start": None, "scada_normal_end": None,
}, {
"pressure_scada_ids": ["J2"],
"scada_normal_start": "2026-10-06T12:00:00+08:00",
"scada_normal_end": "2026-10-06T13:00:00+08:00",
}])
@pytest.mark.parametrize("missing_normal", [False, True])
def test_burst_location_mcp_reuses_project_sensors_and_normal_window_defaults(monkeypatch, tmp_path, overrides, missing_normal):
from app.services import burst_location
args = TOOLS["run_burst_location"].model.model_validate({
"context_handle": HANDLE, "burst_leakage": 0.0274,
"scada_burst_start": START, "scada_burst_end": END, **overrides,
})
devices = [{"device_type": "pressure", "node_id": node, "api_query_id": f"q-{node}"}
for node in ("J1", "J2")]
monkeypatch.setattr(burst_location, "get_all_scada_info", lambda network: devices)
queries = []
def read_scada(**kwargs):
queries.append(kwargs)
if missing_normal and len(queries) == 2:
return {}
return {device: [{"value": 20.0}] for device in kwargs["device_ids"]}
monkeypatch.setattr(burst_location.InternalQueries, "query_scada_by_ids_timerange", read_scada)
@contextmanager
def inp(network, *, purpose):
yield tmp_path / "model.inp"
monkeypatch.setattr(burst_location, "temporary_project_inp", inp)
algorithm = []
def locate(**kwargs):
algorithm.append(kwargs)
return {"located_pipe": "candidate"}
monkeypatch.setattr(burst_location, "run_burst_location", locate)
if missing_normal:
with pytest.raises(ValueError, match="正常压力数据.*无有效数据"):
asyncio.run(analysis_service.calculate(
"run_burst_location", args, network="trusted-project", username="trusted-user",
))
assert algorithm == []
return
result = asyncio.run(analysis_service.calculate(
"run_burst_location", args, network="trusted-project", username="trusted-user",
))
expected_ids = overrides.get("pressure_scada_ids") or ["J1", "J2"]
assert result["pressure_scada_ids"] == expected_ids
assert [q["device_ids"] for q in queries] == [[f"q-{node}" for node in expected_ids]] * 2
normal_start = overrides.get("scada_normal_start") or "2026-10-07T12:00:00+08:00"
normal_end = overrides.get("scada_normal_end") or "2026-10-07T13:00:00+08:00"
assert datetime.fromisoformat(result["scada_window"]["normal_start"]) == datetime.fromisoformat(normal_start)
assert datetime.fromisoformat(result["scada_window"]["normal_end"]) == datetime.fromisoformat(normal_end)
assert algorithm[0]["pressure_scada_ids"] == expected_ids
assert algorithm[0]["final_candidates_csv_path"] is None
@pytest.mark.parametrize("overrides", [
{"scada_normal_start": "2026-10-07T12:00:00+08:00"},
{"scada_normal_end": "2026-10-07T13:00:00+08:00"},
{"scada_normal_start": "2026-10-07T12:00:00", "scada_normal_end": "2026-10-07T13:00:00"},
{"pressure_scada_ids": []}, {"pressure_scada_ids": ["J1", "J1"]},
{"pressure_scada_ids": [str(i) for i in range(21)]},
])
def test_burst_location_optional_inputs_still_reject_partial_or_invalid_values(overrides):
with pytest.raises(ValidationError):
TOOLS["run_burst_location"].model.model_validate({
"context_handle": HANDLE, "burst_leakage": 0.0274,
"scada_burst_start": START, "scada_burst_end": END, **overrides,
})
@pytest.mark.parametrize("name", INPUTS)
def test_calculations_accept_agent_execution_budget_separately_from_duration(name):
args = arguments(name, timeout_seconds=900)
assert args.timeout_seconds == 900
if hasattr(args, "duration"):
assert args.duration == INPUTS[name].get("duration", 900)
@pytest.mark.parametrize("name,module,function", [
("run_burst_detection", "app.services.burst_detection", "run_burst_detection"),
("run_burst_location", "app.services.burst_location", "run_burst_location_by_network"),
("analyze_valve_isolation", "app.services.valve_isolation", "analyze_valve_isolation"),
("run_leakage_identification", "app.services.dma_leakage_estimation", "run_leakage_identification"),
])
def test_adapter_uses_existing_service_without_scheme_or_shared_output(monkeypatch, name, module, function):
calls = []
def calculate(**kwargs):
calls.append(kwargs)
return {"ok": True}
monkeypatch.setattr(f"{module}.{function}", calculate)
result = asyncio.run(analysis_service.calculate(name, arguments(name), network="trusted-project", username="trusted-user"))
assert result == {"ok": True}
call = calls[0]
assert call["network"] == "trusted-project"
assert "context_handle" not in call
if name != "analyze_valve_isolation":
assert call["username"] == "trusted-user"
assert call["scheme_name"] is None
if name == "run_burst_location":
assert call["final_candidates_csv_path"] is None
if name == "run_leakage_identification":
assert call["duration"] == 1
assert call["n_workers"] == 1
@pytest.mark.parametrize("name", ["run_burst_location", "run_simulation"])
def test_dispatch_ends_read_transaction_and_preserves_analysis_provenance(monkeypatch, name):
order = []
async def commit():
order.append("commit")
async def calculate(*args, **kwargs):
order.append("calculate")
assert kwargs == {"network": "trusted-project", "username": "trusted-user"}
return {"located_pipe": "candidate"}
@asynccontextmanager
async def simulation_routing():
yield
monkeypatch.setattr(query_service, "authorize_query", AsyncMock())
monkeypatch.setattr(analysis_service, "calculate", calculate)
scope = SimpleNamespace(connection=SimpleNamespace(commit=commit), project=SimpleNamespace(project_code="trusted-project"), user=SimpleNamespace(username="trusted-user"), simulation_routing=simulation_routing)
result = asyncio.run(query_service.execute_query(scope, object(), TOOLS[name], arguments(name)))
assert order == ["commit", "calculate"]
assert result.semantics == "analysis"
if name == "run_burst_location":
assert result.data["parameters"]["burst_leakage"] == 0.02
else:
assert result.data["parameters"]["duration"] == 3600
assert result.source == "run_simulation"
assert "context_handle" not in result.data["parameters"]
assert result.data["result"]["located_pipe"] == "candidate"
assert "not observed SCADA" in result.data_gaps[-1]
def test_busy_calculation_does_not_start_another_worker(monkeypatch):
from threading import BoundedSemaphore
slots = BoundedSemaphore(1)
slots.acquire()
monkeypatch.setattr(analysis_service, "_slots", slots)
with pytest.raises(HTTPException) as denied:
analysis_service._calculate("analyze_valve_isolation", {}, "project", "user")
assert denied.value.status_code == 503
@pytest.mark.parametrize("name,function", [
("run_contaminant_simulation", "contaminant_simulation"),
("analyze_valve_closure", "valve_close_analysis"),
("analyze_flushing", "flushing_analysis"),
("analyze_pressure_regulation", "pressure_regulation"),
])
def test_scenario_adapters_are_nonpersistent_and_return_only_requested_elements(monkeypatch, name, function):
from app.mcp.scenario_service import calculate_scenario
calls = []
def scenario(**kwargs):
calls.append(kwargs)
return {"node_results": [{"node": "J1", "result": []}, {"node": "J2", "result": []}], "link_results": [{"link": "L1", "result": []}]}
monkeypatch.setattr(f"app.services.simulation_scenarios.{function}", scenario)
monkeypatch.setattr("app.native.wndb.model.options.get_option", lambda _: {"UNITS": "CMH"})
monkeypatch.setattr("app.services.simulation.query_corresponding_element_id_and_query_id", lambda _: SimpleNamespace(fixed_pumps={"P1": "S1"}, variable_pumps={}))
result = calculate_scenario(name, arguments(name).model_dump(exclude={"context_handle"}, exclude_none=True), network="trusted", username="actor")
assert calls[0]["name"] == "trusted"
assert calls[0]["persist_results"] is False
assert calls[0]["scheme_name"] is None
assert result["node_results"] == [{"node": "J1", "result": []}]
assert result["link_results"] == []
assert result["coverage"]["missing_nodes"] == []
if name == "analyze_flushing":
assert calls[0]["flushing_flow"] == 10
def test_sensor_preview_validates_but_does_not_store(monkeypatch):
from app.services import sensor_placement as service
from unittest.mock import Mock
validate = Mock()
save = Mock()
monkeypatch.setattr(service, "validate_sensor_placement_nodes", validate)
monkeypatch.setattr(service, "_sensor_points", lambda *_: [{"node_id": "J1"}])
monkeypatch.setattr(service.sensor_placement_repository, "create_sensor_placement", save)
result = service._create_validated_placement("project", run_name="preview", min_diameter=0,
created_by="actor", sensor_locations=["J1"], persist_results=False)
validate.assert_called_once_with("project", ["J1"])
save.assert_not_called()
assert result["sensor_locations"] == ["J1"]
@pytest.mark.parametrize("name", ["run_contaminant_simulation", "run_simulation"])
def test_scenario_also_requires_scada_permission_even_with_task_read_scope(monkeypatch, name):
permissions = {"simulation.run"}
monkeypatch.setattr(query_service, "permissions_for_context", lambda _: permissions)
binding = SimpleNamespace(tool_cap=list(TOOLS), purpose="task_execution", permission_cap=list(permissions), task_id=uuid4(), task_scope_cap=["scada_readings"])
with pytest.raises(HTTPException):
asyncio.run(query_service.authorize_query(SimpleNamespace(project=object()), binding, TOOLS[name], arguments(name)))
@pytest.mark.parametrize("fails", [False, True])
def test_standard_simulation_uses_online_service_and_propagates_failure(monkeypatch, fails):
from app.mcp.scenario_service import calculate_scenario
events = []
def temporary_project(*_):
pytest.fail("Online calculation must not copy a database")
def simulate(network, start, duration, **kwargs):
assert network == "trusted-project"
assert start.isoformat() == START
assert duration == 3600
assert kwargs == {"output_node_ids": ["J1"], "output_link_ids": ["L1"]}
events.append("calculate")
if fails:
raise RuntimeError("calculation failed")
return {"node_results": [{"node": "J1"}, {"node": "J2"}],
"link_results": [{"link": "L1"}, {"link": "L2"}],
"times": {"report_step": 900}, "units": {"pressure": "m", "flow": "LPS"}}
monkeypatch.setattr("app.native.wndb.core.projects.temporary_project_database", temporary_project)
monkeypatch.setattr("app.services.simulation.run_online_simulation", simulate)
def execute():
return calculate_scenario("run_simulation", arguments("run_simulation").model_dump(
exclude={"context_handle"}), network="trusted-project", username="actor")
if fails:
with pytest.raises(RuntimeError, match="calculation failed"):
execute()
else:
result = execute()
assert result["node_results"] == [{"node": "J1"}]
assert result["link_results"] == [{"link": "L1"}]
assert result["duration_seconds"] == 3600
assert result["persisted_scheme"] is False
assert result["persisted_baseline"] is True
assert result["units"]["pressure"] == "m"
assert events == ["calculate"]