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.
This commit is contained in:
@@ -0,0 +1,139 @@
|
||||
"""Bounded calculation inputs; identity, routing and output paths are host-owned."""
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Annotated, Literal
|
||||
from uuid import UUID
|
||||
|
||||
from pydantic import Field, model_validator
|
||||
|
||||
from app.mcp.models import Id, StrictInput
|
||||
|
||||
|
||||
MAX_CALCULATION_TIMEOUT_SECONDS = 86400
|
||||
|
||||
|
||||
class CalculationInput(StrictInput):
|
||||
timeout_seconds: int = Field(
|
||||
default=120, ge=1, le=MAX_CALCULATION_TIMEOUT_SECONDS, strict=True,
|
||||
description="本次计算等待预算，秒；由Agent按模型规模和任务时长选择，受工具目录中的服务端上限约束；与模拟duration不同，超时不保证已启动的计算停止或回滚",
|
||||
)
|
||||
|
||||
|
||||
class BurstDetectionInput(CalculationInput):
|
||||
target_time: datetime = Field(description="带时区的侦测时刻；读取此前15天同周期压力数据，不是单小时定位")
|
||||
sensor_nodes: list[Id] = Field(min_length=5, max_length=20, description="已核实的压力监测节点ID，至少5个完整传感器")
|
||||
sampling_interval_minutes: int = Field(default=15, ge=1, le=60)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def valid_target(self):
|
||||
if self.target_time.tzinfo is None:
|
||||
raise ValueError("target_time requires explicit timezone")
|
||||
if 1440 % self.sampling_interval_minutes:
|
||||
raise ValueError("sampling interval must divide a day")
|
||||
if len(set(self.sensor_nodes)) != len(self.sensor_nodes):
|
||||
raise ValueError("sensor_nodes must be distinct")
|
||||
return self
|
||||
|
||||
|
||||
class BurstLocationInput(CalculationInput):
|
||||
data_source: Literal["monitoring", "simulation"] = "monitoring"
|
||||
simulation_run_id: UUID | None = Field(default=None, description="已有分析运行UUID，不是方案名称；simulation模式只读取已有异常结果及模型基线")
|
||||
burst_leakage: float = Field(gt=0, allow_inf_nan=False, description="有依据的假定漏水流量，WNTR SI单位m³/s；不是已知爆管管段，不能随意猜测")
|
||||
pressure_scada_ids: list[Id] | None = Field(default=None, min_length=1, max_length=20, description="压力监测设备关联的模型节点ID，不是设备查询编号；省略时自动读取当前项目SCADA配置中的压力测点，不需要猜测节点ID")
|
||||
flow_scada_ids: list[Id] | None = Field(default=None, min_length=1, max_length=20, description="流量监测设备关联的模型要素ID，不是设备查询编号")
|
||||
scada_burst_start: datetime
|
||||
scada_burst_end: datetime
|
||||
scada_normal_start: datetime | None = Field(default=None, description="monitoring模式正常对照窗开始，须与结束时间成对提供；均省略时用爆管窗前一天相同时段，不代表该时段已确认正常")
|
||||
scada_normal_end: datetime | None = Field(default=None, description="monitoring模式正常对照窗结束，须与开始时间成对提供；simulation模式不填写，使用同时间窗模型基线")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def valid_windows(self):
|
||||
pairs = [(self.scada_burst_start, self.scada_burst_end)]
|
||||
if self.data_source == "monitoring":
|
||||
if self.simulation_run_id is not None:
|
||||
raise ValueError("monitoring cannot use simulation_run_id")
|
||||
if self.scada_normal_start is not None or self.scada_normal_end is not None:
|
||||
pairs.append((self.scada_normal_start, self.scada_normal_end))
|
||||
elif self.simulation_run_id is None or self.scada_normal_start is not None or self.scada_normal_end is not None:
|
||||
raise ValueError("simulation requires run_id and uses the same-window model baseline")
|
||||
for start, end in pairs:
|
||||
if start is None or end is None or start.tzinfo is None or end.tzinfo is None:
|
||||
raise ValueError("explicit timezone-aware burst and normal windows are required")
|
||||
if not timedelta(0) < end - start <= timedelta(days=1):
|
||||
raise ValueError("each window must be positive and at most 24 hours")
|
||||
for ids in (self.pressure_scada_ids, self.flow_scada_ids):
|
||||
if ids is not None and len(set(ids)) != len(ids):
|
||||
raise ValueError("sensor IDs must be distinct")
|
||||
return self
|
||||
|
||||
|
||||
class ValveIsolationInput(CalculationInput):
|
||||
accident_element: list[Id] = Field(min_length=1, max_length=20, description="经核实或明确作为候选的事故管段/节点ID")
|
||||
disabled_valves: list[Id] = Field(default_factory=list, max_length=200, description="已知无法关闭的阀门；未提供不代表现场确认全部可用")
|
||||
|
||||
|
||||
class LeakageIdentificationInput(CalculationInput):
|
||||
scada_start: datetime
|
||||
scada_end: datetime
|
||||
sensor_nodes: list[Id] = Field(min_length=1, max_length=20)
|
||||
q_sum: float = Field(gt=0, allow_inf_nan=False, description="有依据的总漏损流量约束，m³/s，不得使用无依据默认值")
|
||||
start_time: float = Field(ge=0, lt=24, allow_inf_nan=False, description="模型起算小时；需与SCADA时间窗及模型时钟对齐")
|
||||
timestep: int = Field(default=15, ge=5, le=60, description="计算步长，分钟；需与观测采样对应")
|
||||
dma_count: int = Field(ge=1, le=20)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def valid_window(self):
|
||||
if self.scada_start.tzinfo is None or self.scada_end.tzinfo is None:
|
||||
raise ValueError("timestamps require explicit timezone")
|
||||
duration = self.scada_end - self.scada_start
|
||||
if not timedelta(0) < duration <= timedelta(days=1):
|
||||
raise ValueError("window must be positive and at most 24 hours")
|
||||
if duration.total_seconds() % (self.timestep * 60):
|
||||
raise ValueError("window must contain whole calculation steps")
|
||||
if len(set(self.sensor_nodes)) != len(self.sensor_nodes):
|
||||
raise ValueError("sensor_nodes must be distinct")
|
||||
return self
|
||||
|
||||
|
||||
class ScenarioInput(CalculationInput):
|
||||
start_time: datetime
|
||||
duration: int = Field(default=900, ge=60, le=86400, description="模拟时长，秒；首次优先短时间窗")
|
||||
output_node_ids: list[Id] = Field(min_length=1, max_length=20, description="需要返回结果的节点；不把局部结果解释为全网无影响")
|
||||
output_link_ids: list[Id] = Field(default_factory=list, max_length=20)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def timezone_required(self):
|
||||
if self.start_time.tzinfo is None:
|
||||
raise ValueError("start_time requires explicit timezone")
|
||||
return self
|
||||
|
||||
|
||||
class ContaminantSimulationInput(ScenarioInput):
|
||||
source: Id = Field(description="假设污染源节点；这是正向传播模拟，不是自动反演污染源")
|
||||
concentration: float = Field(gt=0, allow_inf_nan=False, description="假设源浓度，mg/L")
|
||||
pattern: Id | None = Field(default=None, description="已存在的污染源时间模式ID；不填为持续恒定源")
|
||||
|
||||
|
||||
class ValveClosureInput(ScenarioInput):
|
||||
valves: list[Id] = Field(min_length=1, max_length=20, description="在临时模型中关闭的阀门，不下发现场控制")
|
||||
|
||||
|
||||
class FlushingInput(ScenarioInput):
|
||||
drainage_node_id: Id
|
||||
flushing_flow_m3h: float = Field(gt=0, allow_inf_nan=False, description="明确的冲洗流量，m³/h；不使用无依据的自动排放量")
|
||||
closed_valves: list[Id] = Field(default_factory=list, max_length=20)
|
||||
|
||||
|
||||
NonnegativeValue = Annotated[float, Field(ge=0, allow_inf_nan=False)]
|
||||
PumpSeries = Annotated[list[NonnegativeValue], Field(min_length=1, max_length=96)]
|
||||
|
||||
|
||||
class PressureRegulationInput(ScenarioInput):
|
||||
pump_control: dict[Id, PumpSeries] = Field(min_length=1, max_length=20, description="已登记泵ID及运行序列；工频值为0/1，变频值为频率Hz，每项对应模型水力步长，不自动优化")
|
||||
tank_init_level: dict[Id, Annotated[float, Field(gt=0, allow_inf_nan=False)]] = Field(default_factory=dict, max_length=20, description="水箱初始水位，须在模型允许范围内")
|
||||
|
||||
|
||||
class SensorPlacementInput(CalculationInput):
|
||||
method: Literal["sensitivity", "kmeans"]
|
||||
sensor_count: int = Field(ge=1, le=20)
|
||||
min_diameter: int = Field(default=0, ge=0, le=5000, description="最小管径，mm")
|
||||
@@ -0,0 +1,69 @@
|
||||
"""Side-effect-limited adapters for existing domain calculations."""
|
||||
|
||||
import asyncio
|
||||
from threading import BoundedSemaphore
|
||||
from typing import Any
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel
|
||||
from app.mcp.scenario_service import SCENARIO_TOOLS, calculate_scenario
|
||||
|
||||
# These services synchronously read SCADA/model series or write online results.
|
||||
# Topology-only isolation and model-based sensor placement need only BizDB.
|
||||
TIMESCALE_CALCULATIONS = SCENARIO_TOOLS | frozenset({
|
||||
"run_burst_detection", "run_burst_location", "run_leakage_identification",
|
||||
})
|
||||
|
||||
# A timed-out HTTP request cannot kill a native algorithm thread. Keep its slot
|
||||
# until the calculation really exits, so retries cannot accumulate CPU jobs.
|
||||
_slots = BoundedSemaphore(2)
|
||||
|
||||
|
||||
def _calculate(name: str, parameters: dict[str, Any], network: str, username: str) -> dict[str, Any]:
|
||||
if not _slots.acquire(blocking=False):
|
||||
raise HTTPException(503, "Calculation capacity is occupied")
|
||||
try:
|
||||
if name in SCENARIO_TOOLS:
|
||||
return calculate_scenario(name, parameters, network=network, username=username)
|
||||
if name == "optimize_sensor_placement":
|
||||
from app.services.sensor_placement import optimize_sensor_placement_by_kmeans, optimize_sensor_placement_by_sensitivity
|
||||
parameters = dict(parameters)
|
||||
method = parameters.pop("method")
|
||||
optimizer = optimize_sensor_placement_by_sensitivity if method == "sensitivity" else optimize_sensor_placement_by_kmeans
|
||||
return optimizer(project_code=network, run_name="MCP preview", created_by=username,
|
||||
persist_results=False, **parameters)
|
||||
if name == "run_leakage_identification":
|
||||
from app.services.dma_leakage_estimation import run_leakage_identification
|
||||
|
||||
duration = (parameters["scada_end"] - parameters["scada_start"]).total_seconds() / 3600
|
||||
return run_leakage_identification(
|
||||
network=network, username=username, scheme_name=None,
|
||||
duration=duration, q_sum_unit="m3/s", output_flow_unit="m3/s",
|
||||
n_workers=1, pop_size=50, max_gen=100, **parameters,
|
||||
)
|
||||
if name == "run_burst_detection":
|
||||
from app.services.burst_detection import run_burst_detection
|
||||
|
||||
return run_burst_detection(network=network, username=username, scheme_name=None, **parameters)
|
||||
if name == "run_burst_location":
|
||||
from app.services.burst_location import run_burst_location_by_network
|
||||
|
||||
return run_burst_location_by_network(
|
||||
network=network, username=username, scheme_name=None,
|
||||
final_candidates_csv_path=None, **parameters,
|
||||
)
|
||||
if name == "analyze_valve_isolation":
|
||||
from app.services.valve_isolation import analyze_valve_isolation
|
||||
|
||||
return analyze_valve_isolation(network=network, **parameters)
|
||||
raise ValueError("Calculation is not registered")
|
||||
finally:
|
||||
_slots.release()
|
||||
|
||||
|
||||
async def calculate(name: str, args: BaseModel, *, network: str, username: str) -> dict[str, Any]:
|
||||
# to_thread copies the validated project's routing context. Never accept
|
||||
# database routes, actor identity, arbitrary measurements or paths from MCP.
|
||||
return await asyncio.to_thread(
|
||||
_calculate, name, args.model_dump(exclude={"context_handle", "timeout_seconds"}, exclude_none=True), network, username,
|
||||
)
|
||||
@@ -0,0 +1,505 @@
|
||||
"""Official SDK transport with product authorization, persistence and query handlers."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from urllib.parse import urlsplit
|
||||
from uuid import uuid5, NAMESPACE_URL
|
||||
from typing import Any
|
||||
|
||||
from fastapi import HTTPException
|
||||
from mcp import types
|
||||
from mcp.server import Server
|
||||
from mcp.server.auth.middleware.auth_context import get_access_token
|
||||
from mcp.server.auth.settings import AuthSettings
|
||||
from mcp.server.auth.routes import create_protected_resource_routes
|
||||
from mcp.server.caching import CacheHint
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from psycopg.errors import UndefinedTable, InvalidSchemaName
|
||||
from pydantic import ValidationError
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
from app.mcp.auth import (
|
||||
McpTokenVerifier,
|
||||
QueryScopeMiddleware,
|
||||
QUERY_SCOPE,
|
||||
CONTEXT_SCOPE,
|
||||
challenge,
|
||||
)
|
||||
from app.mcp.catalog import TOOLS
|
||||
from app.mcp.config import McpSettings
|
||||
from app.mcp.contexts import resolve_context, verify_task
|
||||
from app.mcp.models import QueryResult
|
||||
from app.mcp.scope import project_scope
|
||||
from app.mcp.query_service import execute_query, task_service, authorize_query
|
||||
from app.infra.db.postgresql.mcp_queries import McpQueryRepository
|
||||
from app.domain.schemas.operational_data import TaskDataSnapshot
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def tool_error(code, message):
|
||||
return types.CallToolResult(
|
||||
isError=True,
|
||||
content=[types.TextContent(type="text", text=message)],
|
||||
structuredContent={"error": {"code": code, "message": message}},
|
||||
)
|
||||
|
||||
|
||||
class ModernOnlyMiddleware:
|
||||
def __init__(self, app):
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] != "http" or scope.get("path") != "/mcp":
|
||||
return await self.app(scope, receive, send)
|
||||
if scope["method"] in {"GET", "DELETE"}:
|
||||
return await JSONResponse(
|
||||
{"error": "modern_protocol_required"},
|
||||
status_code=405,
|
||||
headers={"Allow": "POST"},
|
||||
)(scope, receive, send)
|
||||
headers = {
|
||||
k.decode("latin1").lower(): v.decode("latin1")
|
||||
for k, v in scope.get("headers", [])
|
||||
}
|
||||
version = headers.get("mcp-protocol-version")
|
||||
if version != "2026-07-28":
|
||||
body = bytearray()
|
||||
request_id = None
|
||||
body_version = None
|
||||
while True:
|
||||
message = await receive()
|
||||
if message["type"] == "http.disconnect":
|
||||
return
|
||||
body.extend(message.get("body", b""))
|
||||
if len(body) > 1024 * 1024:
|
||||
return await JSONResponse(
|
||||
{"error": "request_too_large"}, status_code=413
|
||||
)(scope, receive, send)
|
||||
if not message.get("more_body"):
|
||||
break
|
||||
try:
|
||||
payload = json.loads(body)
|
||||
request_id = payload.get("id")
|
||||
body_version = payload.get("params", {}).get("_meta", {}).get(
|
||||
types.PROTOCOL_VERSION_META_KEY
|
||||
)
|
||||
except (ValueError, AttributeError):
|
||||
pass
|
||||
mismatch = not version or (
|
||||
isinstance(body_version, str) and body_version != version
|
||||
)
|
||||
code = -32020 if mismatch else -32022
|
||||
result = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"error": {
|
||||
"code": code,
|
||||
"message": (
|
||||
"MCP-Protocol-Version header is missing or differs from request metadata"
|
||||
if mismatch else "MCP 2026-07-28 is required"
|
||||
),
|
||||
"data": {"supported": ["2026-07-28"], "requested": version},
|
||||
},
|
||||
}
|
||||
return await JSONResponse(result, status_code=400)(scope, receive, send)
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
|
||||
class McpQueryApplication:
|
||||
def __init__(
|
||||
self,
|
||||
settings: McpSettings,
|
||||
*,
|
||||
verifier=None,
|
||||
scope_factory=project_scope,
|
||||
executor=execute_query,
|
||||
):
|
||||
self.settings = settings
|
||||
self.verifier = verifier or McpTokenVerifier(settings)
|
||||
self.scope_factory = scope_factory
|
||||
self.executor = executor
|
||||
self.active = {}
|
||||
self.ready = False
|
||||
self.server = Server(
|
||||
"tjwater-query",
|
||||
version="1.0.0",
|
||||
instructions="Project-bound water queries. Model baselines are not observations. Truncated results do not cover the full range.",
|
||||
cache_hints={
|
||||
"server/discover": CacheHint(300000, "public"),
|
||||
"tools/list": CacheHint(300000, "public"),
|
||||
},
|
||||
on_list_tools=self.list_tools,
|
||||
on_call_tool=self.call_tool,
|
||||
get_tool_input_schema=self.tool_input_schema,
|
||||
on_list_resources=self.list_resources,
|
||||
on_list_resource_templates=self.list_templates,
|
||||
on_read_resource=self.read_resource,
|
||||
)
|
||||
self.server.add_request_handler(
|
||||
"server/discover", types.RequestParams, self.discover
|
||||
)
|
||||
app = self.server.streamable_http_app(
|
||||
stateless_http=True,
|
||||
json_response=True,
|
||||
max_request_body_size=1024 * 1024,
|
||||
auth=AuthSettings(
|
||||
issuer_url=settings.issuer_uri,
|
||||
resource_server_url=settings.resource_uri,
|
||||
required_scopes=["mcp.catalog.read"],
|
||||
validate_token_resource=True,
|
||||
),
|
||||
token_verifier=self.verifier,
|
||||
transport_security=TransportSecuritySettings(
|
||||
allowed_hosts=settings.allowed_hosts,
|
||||
allowed_origins=settings.allowed_origins,
|
||||
),
|
||||
)
|
||||
# Catalog authentication and advertised resource scopes are different concerns.
|
||||
metadata = create_protected_resource_routes(
|
||||
resource_url=settings.resource_uri,
|
||||
authorization_servers=[settings.issuer_uri],
|
||||
scopes_supported=["mcp.catalog.read", QUERY_SCOPE, CONTEXT_SCOPE],
|
||||
resource_name="TJWater Query",
|
||||
)
|
||||
metadata_path = metadata[0].path
|
||||
app.router.routes = [
|
||||
route for route in app.router.routes if route.path != metadata_path
|
||||
] + metadata
|
||||
self.asgi = ModernOnlyMiddleware(
|
||||
QueryScopeMiddleware(app, self.verifier, settings)
|
||||
)
|
||||
from app.api.v1.endpoints.mcp_queries import router
|
||||
|
||||
self.router = router
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(self):
|
||||
async with self.server.session_manager.run():
|
||||
self.ready = True
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.ready = False
|
||||
|
||||
async def discover(self, ctx, params):
|
||||
return types.DiscoverResult(
|
||||
supported_versions=["2026-07-28"],
|
||||
capabilities=self.server.get_capabilities(
|
||||
protocol_version=ctx.protocol_version,
|
||||
extensions={"io.modelcontextprotocol/oauth-client-credentials": {}},
|
||||
),
|
||||
instructions=self.server.instructions,
|
||||
)
|
||||
|
||||
def tool_input_schema(self, name: str) -> dict[str, Any] | None:
|
||||
tool = TOOLS.get(name)
|
||||
if tool is None:
|
||||
return None
|
||||
schema = tool.model.model_json_schema()
|
||||
if tool.computation:
|
||||
schema["properties"]["timeout_seconds"]["maximum"] = self.settings.max_calculation_timeout_seconds
|
||||
return schema
|
||||
|
||||
async def list_tools(self, ctx, params):
|
||||
if params and params.cursor:
|
||||
raise MCPError(-32602, "The query catalog has no continuation page")
|
||||
tools = []
|
||||
for name in sorted(TOOLS):
|
||||
tool = TOOLS[name].protocol_tool()
|
||||
tool.input_schema.update(self.tool_input_schema(name))
|
||||
tools.append(tool)
|
||||
return types.ListToolsResult(
|
||||
tools=tools,
|
||||
ttlMs=300000,
|
||||
cacheScope="public",
|
||||
)
|
||||
|
||||
async def list_resources(self, ctx, params):
|
||||
return types.ListResourcesResult(resources=[], ttlMs=0, cacheScope="private")
|
||||
|
||||
async def list_templates(self, ctx, params):
|
||||
return types.ListResourceTemplatesResult(
|
||||
resourceTemplates=[
|
||||
types.ResourceTemplate(
|
||||
name="task-evidence",
|
||||
uriTemplate="tjwater://tasks/{task_id}/evidence/{evidence_id}",
|
||||
mimeType="application/json",
|
||||
)
|
||||
],
|
||||
ttlMs=300000,
|
||||
cacheScope="public",
|
||||
)
|
||||
|
||||
async def read_resource(self, ctx, params):
|
||||
meta = ctx.meta or {}
|
||||
handle = meta.get("io.tjwater/contextHandle")
|
||||
url = urlsplit(str(params.uri))
|
||||
parts = url.path.strip("/").split("/")
|
||||
if (
|
||||
url.scheme != "tjwater"
|
||||
or url.netloc != "tasks"
|
||||
or len(parts) != 3
|
||||
or parts[1] != "evidence"
|
||||
or url.query
|
||||
or url.fragment
|
||||
):
|
||||
raise MCPError(-32602, "Unknown query resource")
|
||||
result = await self._query(
|
||||
ctx,
|
||||
"get_task_evidence",
|
||||
{"context_handle": handle, "evidence_id": parts[2]},
|
||||
resource_task=parts[0],
|
||||
include_full=True,
|
||||
)
|
||||
if result.isError:
|
||||
raise MCPError(
|
||||
-32602, "Evidence is unavailable in this authorization context"
|
||||
)
|
||||
return types.ReadResourceResult(
|
||||
contents=[
|
||||
types.TextResourceContents(
|
||||
uri=str(params.uri),
|
||||
mimeType="application/json",
|
||||
text=json.dumps(result.structuredContent, ensure_ascii=False),
|
||||
)
|
||||
],
|
||||
ttlMs=0,
|
||||
cacheScope="private",
|
||||
)
|
||||
|
||||
async def call_tool(self, ctx, params):
|
||||
return await self._query(ctx, params.name, params.arguments or {})
|
||||
|
||||
async def _query(
|
||||
self, ctx, name, arguments, resource_task=None, include_full=False
|
||||
):
|
||||
tool = TOOLS.get(name)
|
||||
if tool is None:
|
||||
raise MCPError(-32602, "Unknown query tool")
|
||||
principal = get_access_token()
|
||||
if principal is None or QUERY_SCOPE not in principal.scopes:
|
||||
return tool_error("PERMISSION_DENIED", "Query permission is required")
|
||||
try:
|
||||
args = tool.model.model_validate(arguments)
|
||||
except ValidationError:
|
||||
return tool_error(
|
||||
"INVALID_QUERY", "Query parameters do not match the tool schema"
|
||||
)
|
||||
timeout_seconds = (
|
||||
args.timeout_seconds if tool.computation else self.settings.query_timeout_seconds
|
||||
)
|
||||
if tool.computation and timeout_seconds > self.settings.max_calculation_timeout_seconds:
|
||||
return tool_error("INVALID_QUERY", "Calculation timeout exceeds the server limit")
|
||||
request = ctx.request
|
||||
project = request.headers.get("x-project-id", "") if request else ""
|
||||
meta = ctx.meta or {}
|
||||
call_id = meta.get("io.tjwater/callId")
|
||||
if not isinstance(call_id, str) or not 0 < len(call_id) <= 128:
|
||||
return tool_error(
|
||||
"CALL_ID_REQUIRED", "A host-owned call identifier is required"
|
||||
)
|
||||
key = (principal.subject, project)
|
||||
if self.active.get(key, 0) >= 4:
|
||||
return tool_error(
|
||||
"RATE_LIMITED", "At most four queries may run concurrently"
|
||||
)
|
||||
self.active[key] = self.active.get(key, 0) + 1
|
||||
try:
|
||||
async with (
|
||||
asyncio.timeout(timeout_seconds),
|
||||
self.scope_factory(principal, project) as scope,
|
||||
):
|
||||
binding = await resolve_context(scope, principal, args.context_handle)
|
||||
if resource_task is not None and str(binding.task_id) != resource_task:
|
||||
raise HTTPException(403, "Resource task mismatch")
|
||||
await authorize_query(scope, binding, tool, args)
|
||||
repo = McpQueryRepository(scope.connection, scope.project.project_id)
|
||||
parameters = {
|
||||
"tool": name,
|
||||
"arguments": args.model_dump(
|
||||
mode="json", exclude={"context_handle"}
|
||||
),
|
||||
}
|
||||
scope_key = json.dumps(
|
||||
[
|
||||
binding.subject,
|
||||
binding.client_id,
|
||||
binding.native_session_id,
|
||||
str(binding.task_id),
|
||||
str(binding.run_id),
|
||||
str(binding.execution_request_id),
|
||||
],
|
||||
separators=(",", ":"),
|
||||
)
|
||||
previous = (
|
||||
await repo.receipt(scope_key, call_id) if binding.task_id else None
|
||||
)
|
||||
if previous:
|
||||
if previous["parameters"] != parameters:
|
||||
raise HTTPException(
|
||||
409, "Call identifier has different query parameters"
|
||||
)
|
||||
result = QueryResult.model_validate(previous["result"])
|
||||
else:
|
||||
result = (
|
||||
await self.executor(
|
||||
scope, binding, tool, args, full_snapshot=True
|
||||
)
|
||||
if include_full
|
||||
else await self.executor(scope, binding, tool, args)
|
||||
)
|
||||
# Never serialize an unbounded response into a model message or database receipt.
|
||||
if len(result.model_dump_json().encode()) > 8 * 1024 * 1024:
|
||||
return tool_error(
|
||||
"RESULT_TOO_LARGE",
|
||||
"Narrow the query or request a smaller page",
|
||||
)
|
||||
if binding.task_id and name not in {
|
||||
"get_task_context",
|
||||
"get_task_evidence",
|
||||
}:
|
||||
async with scope.connection.transaction():
|
||||
await repo.lock_receipt(scope_key, call_id)
|
||||
previous = await repo.receipt(scope_key, call_id)
|
||||
if previous:
|
||||
if previous["parameters"] != parameters:
|
||||
raise HTTPException(
|
||||
409,
|
||||
"Call identifier has different query parameters",
|
||||
)
|
||||
result = QueryResult.model_validate(previous["result"])
|
||||
else:
|
||||
task_repo = task_service(scope).repo
|
||||
# Lock the task row only after external/data-source work has finished.
|
||||
await task_repo.lock(f"task:{binding.task_id}")
|
||||
await verify_task(scope, binding)
|
||||
result.source_ref = "mcp:" + str(
|
||||
uuid5(NAMESPACE_URL, scope_key + ":" + call_id)
|
||||
)
|
||||
snapshot = TaskDataSnapshot(
|
||||
source=name,
|
||||
source_ref=result.source_ref,
|
||||
queried_at=result.queried_at,
|
||||
origin=result.origin,
|
||||
data=result.model_dump(mode="json"),
|
||||
truncated=result.truncated,
|
||||
)
|
||||
evidence, _ = await task_service(
|
||||
scope
|
||||
).record_data_snapshot(binding.task_id, snapshot)
|
||||
result.evidence_id = evidence.evidence_id
|
||||
await repo.save_receipt(
|
||||
scope_key,
|
||||
call_id,
|
||||
parameters,
|
||||
result.model_dump(mode="json"),
|
||||
)
|
||||
output = result.model_dump(mode="json")
|
||||
content = []
|
||||
if (
|
||||
not include_full
|
||||
and len(json.dumps(output, ensure_ascii=False).encode()) > 12000
|
||||
):
|
||||
evidence_id = result.evidence_id or (
|
||||
args.evidence_id if name == "get_task_evidence" else None
|
||||
)
|
||||
if binding.task_id and evidence_id:
|
||||
uri = (
|
||||
f"tjwater://tasks/{binding.task_id}/evidence/{evidence_id}"
|
||||
)
|
||||
output["data"] = {
|
||||
"resource_uri": uri,
|
||||
"preview": "The complete query snapshot is saved in task evidence",
|
||||
}
|
||||
output["truncated"] = True
|
||||
output["data_gaps"] = [
|
||||
*result.data_gaps,
|
||||
"Inline result abbreviated; read the saved evidence resource",
|
||||
]
|
||||
content.append(
|
||||
types.ResourceLink(
|
||||
type="resource_link",
|
||||
uri=uri,
|
||||
name="query-evidence",
|
||||
mimeType="application/json",
|
||||
)
|
||||
)
|
||||
else:
|
||||
return tool_error(
|
||||
"RESULT_TOO_LARGE",
|
||||
"Request a smaller page or narrower window",
|
||||
)
|
||||
content.insert(
|
||||
0,
|
||||
types.TextContent(
|
||||
type="text", text=json.dumps(output, ensure_ascii=False)
|
||||
),
|
||||
)
|
||||
return types.CallToolResult(content=content, structuredContent=output)
|
||||
except HTTPException as exc:
|
||||
codes = {
|
||||
401: "AUTH_REQUIRED",
|
||||
403: "QUERY_CONTEXT_DENIED",
|
||||
404: "NOT_FOUND",
|
||||
409: "QUERY_CONTEXT_STALE",
|
||||
413: "RESULT_TOO_LARGE",
|
||||
422: "INVALID_QUERY",
|
||||
503: "SOURCE_UNAVAILABLE",
|
||||
}
|
||||
return tool_error(
|
||||
codes.get(exc.status_code, "QUERY_FAILED"),
|
||||
str(exc.detail)
|
||||
if isinstance(exc.detail, str)
|
||||
and exc.status_code in {403, 404, 409, 413, 422}
|
||||
else "Query is unavailable",
|
||||
)
|
||||
except (UndefinedTable, InvalidSchemaName):
|
||||
return tool_error(
|
||||
"MCP_STORAGE_UNAVAILABLE",
|
||||
"MCP project storage has not been provisioned",
|
||||
)
|
||||
except TimeoutError:
|
||||
return tool_error("QUERY_TIMEOUT", "Query execution timed out")
|
||||
except ValueError:
|
||||
return tool_error(
|
||||
"INVALID_QUERY", "Invalid query value or incompatible source data"
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"MCP query unavailable", extra={"tool": name, "project_id": project}
|
||||
)
|
||||
return tool_error("SOURCE_UNAVAILABLE", "Query source is unavailable")
|
||||
finally:
|
||||
remaining = self.active[key] - 1
|
||||
if remaining:
|
||||
self.active[key] = remaining
|
||||
else:
|
||||
self.active.pop(key, None)
|
||||
|
||||
async def identity(self, request):
|
||||
auth = request.headers.get("authorization", "")
|
||||
try:
|
||||
principal = (
|
||||
await self.verifier.verify_token(auth[7:])
|
||||
if auth.lower().startswith("bearer ")
|
||||
else None
|
||||
)
|
||||
except Exception:
|
||||
raise HTTPException(503, "OAuth key verification is unavailable") from None
|
||||
if not principal:
|
||||
raise HTTPException(
|
||||
401,
|
||||
"MCP authentication required",
|
||||
headers=challenge(self.settings, CONTEXT_SCOPE),
|
||||
)
|
||||
if CONTEXT_SCOPE not in principal.scopes:
|
||||
raise HTTPException(
|
||||
403,
|
||||
"MCP context scope required",
|
||||
headers=challenge(self.settings, CONTEXT_SCOPE, insufficient=True),
|
||||
)
|
||||
return principal
|
||||
+169
@@ -0,0 +1,169 @@
|
||||
"""Verify resource-bound Keycloak tokens without passing them to business APIs."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from uuid import UUID
|
||||
|
||||
import httpx
|
||||
from jose import JWTError, jwt
|
||||
from mcp.server.auth.provider import AccessToken
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
from app.mcp.config import McpSettings
|
||||
|
||||
CATALOG_SCOPE = "mcp.catalog.read"
|
||||
QUERY_SCOPE = "mcp.query"
|
||||
CONTEXT_SCOPE = "mcp.context.bind"
|
||||
|
||||
|
||||
class McpTokenVerifier:
|
||||
def __init__(self, settings: McpSettings, client: httpx.AsyncClient | None = None):
|
||||
self.settings = settings
|
||||
self.client = client
|
||||
self._keys: list[dict] = []
|
||||
self._expires = 0.0
|
||||
self._last_refresh = 0.0
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def _load_keys(self, kid: str) -> list[dict]:
|
||||
now = time.monotonic()
|
||||
if self._expires > now and any(k.get("kid") == kid for k in self._keys):
|
||||
return self._keys
|
||||
async with self._lock:
|
||||
now = time.monotonic()
|
||||
if self._expires > now and any(k.get("kid") == kid for k in self._keys):
|
||||
return self._keys
|
||||
if self._expires > now and now - self._last_refresh < 20:
|
||||
return self._keys
|
||||
url = (
|
||||
self.settings.issuer_uri.rstrip("/") + "/protocol/openid-connect/certs"
|
||||
)
|
||||
self._last_refresh = now
|
||||
if self.client:
|
||||
response = await self.client.get(url)
|
||||
else:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=5, trust_env=False, follow_redirects=False
|
||||
) as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError:
|
||||
raise httpx.RequestError("invalid OAuth key set response") from None
|
||||
keys = payload.get("keys") if isinstance(payload, dict) else None
|
||||
if not isinstance(keys, list) or len(keys) > 100:
|
||||
raise httpx.RequestError("invalid OAuth key set")
|
||||
self._keys = [
|
||||
k
|
||||
for k in keys
|
||||
if isinstance(k, dict)
|
||||
and k.get("kty") == "RSA"
|
||||
and k.get("use", "sig") == "sig"
|
||||
and k.get("alg", "RS256") == "RS256"
|
||||
]
|
||||
self._expires = now + 300
|
||||
return self._keys
|
||||
|
||||
async def verify_token(self, token: str) -> AccessToken | None:
|
||||
if not token or len(token) > 16384:
|
||||
return None
|
||||
try:
|
||||
header = jwt.get_unverified_header(token)
|
||||
kid = header.get("kid")
|
||||
if (
|
||||
header.get("alg") != "RS256"
|
||||
or not isinstance(kid, str)
|
||||
or not 0 < len(kid) <= 128
|
||||
):
|
||||
return None
|
||||
keys = await self._load_keys(kid)
|
||||
key = next((k for k in keys if k.get("kid") == kid), None)
|
||||
if key is None:
|
||||
return None
|
||||
claims = jwt.decode(
|
||||
token,
|
||||
key,
|
||||
algorithms=["RS256"],
|
||||
audience=self.settings.audience,
|
||||
issuer=self.settings.issuer_uri,
|
||||
options={"require_exp": True, "require_iat": True, "require_sub": True},
|
||||
)
|
||||
UUID(claims["sub"])
|
||||
client_id = claims.get("azp")
|
||||
if not isinstance(client_id, str) or not client_id:
|
||||
return None
|
||||
if (
|
||||
not isinstance(claims.get("iat"), (int, float))
|
||||
or not time.time() - 900 < claims["iat"] <= time.time() + 30
|
||||
or claims["exp"] > claims["iat"] + 900
|
||||
):
|
||||
return None
|
||||
return AccessToken(
|
||||
token=token,
|
||||
client_id=client_id,
|
||||
subject=claims["sub"],
|
||||
scopes=str(claims.get("scope", "")).split(),
|
||||
expires_at=int(claims["exp"]),
|
||||
resource=self.settings.resource_uri,
|
||||
claims=claims,
|
||||
)
|
||||
except (JWTError, ValueError, KeyError, TypeError):
|
||||
return None
|
||||
# JWKS/network failures propagate as unavailable, never as empty data.
|
||||
|
||||
|
||||
def challenge(
|
||||
settings: McpSettings, scope: str, *, insufficient: bool = False
|
||||
) -> dict[str, str]:
|
||||
from mcp.server.auth.routes import build_resource_metadata_url
|
||||
|
||||
metadata = build_resource_metadata_url(settings.resource_uri)
|
||||
suffix = ', error="insufficient_scope"' if insufficient else ""
|
||||
return {
|
||||
"WWW-Authenticate": f'Bearer resource_metadata="{metadata}", scope="{scope}"{suffix}'
|
||||
}
|
||||
|
||||
|
||||
class QueryScopeMiddleware:
|
||||
"""HTTP scope challenges before JSON-RPC dispatch; SDK validates mirrored headers."""
|
||||
|
||||
def __init__(self, app, verifier: McpTokenVerifier, settings: McpSettings):
|
||||
self.app, self.verifier, self.settings = app, verifier, settings
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] != "http" or scope.get("path") != "/mcp":
|
||||
return await self.app(scope, receive, send)
|
||||
headers = {
|
||||
k.decode("latin1").lower(): v.decode("latin1")
|
||||
for k, v in scope.get("headers", [])
|
||||
}
|
||||
auth = headers.get("authorization", "")
|
||||
try:
|
||||
token = (
|
||||
await self.verifier.verify_token(auth[7:])
|
||||
if auth.lower().startswith("bearer ")
|
||||
else None
|
||||
)
|
||||
except (httpx.HTTPError, ValueError):
|
||||
return await JSONResponse(
|
||||
{"error": "authentication_unavailable"}, status_code=503
|
||||
)(scope, receive, send)
|
||||
required = (
|
||||
QUERY_SCOPE
|
||||
if headers.get("mcp-method") in {"tools/call", "resources/read"}
|
||||
else CATALOG_SCOPE
|
||||
)
|
||||
if token is None:
|
||||
return await JSONResponse(
|
||||
{"error": "invalid_token"},
|
||||
status_code=401,
|
||||
headers=challenge(self.settings, required),
|
||||
)(scope, receive, send)
|
||||
if required not in token.scopes:
|
||||
return await JSONResponse(
|
||||
{"error": "insufficient_scope"},
|
||||
status_code=403,
|
||||
headers=challenge(self.settings, required, insufficient=True),
|
||||
)(scope, receive, send)
|
||||
await self.app(scope, receive, send)
|
||||
@@ -0,0 +1,221 @@
|
||||
"""One finite query catalog; names, schemas and permissions come from this registry."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pydantic import BaseModel
|
||||
from mcp import types
|
||||
from app.mcp import models as m
|
||||
from app.mcp.analysis_models import BurstDetectionInput, BurstLocationInput, ValveIsolationInput, LeakageIdentificationInput
|
||||
from app.mcp.analysis_models import ScenarioInput, ContaminantSimulationInput, ValveClosureInput, FlushingInput, PressureRegulationInput, SensorPlacementInput
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueryTool:
|
||||
name: str
|
||||
description: str
|
||||
model: type[BaseModel]
|
||||
permission: str
|
||||
semantics: str
|
||||
task_only: bool = False
|
||||
computation: bool = False
|
||||
|
||||
def protocol_tool(self) -> types.Tool:
|
||||
if self.name == "run_simulation":
|
||||
suffix = "。任务调用复用证据留存。"
|
||||
elif self.computation:
|
||||
suffix = "。仅计算建议，不保存方案、不控制设备；任务调用复用证据留存。"
|
||||
elif self.name == "query_nearby_elements":
|
||||
suffix = "。仅查询当前已授权项目的模型管网，按距离和元素ID稳定排序；按需分页，不自动遍历全量；空间邻近不代表管网连通或供水关系。"
|
||||
else:
|
||||
suffix = (
|
||||
"。先限定对象和时间范围，再读取明细；集合默认20条，按需逐页扩大，"
|
||||
"不要自动遍历全量；截断或非零offset不能证明完整范围无异常。"
|
||||
)
|
||||
return types.Tool(
|
||||
name=self.name,
|
||||
description=self.description + suffix,
|
||||
inputSchema=self.model.model_json_schema(),
|
||||
outputSchema=m.QueryResult.model_json_schema(),
|
||||
annotations=types.ToolAnnotations(
|
||||
readOnlyHint=False, destructiveHint=self.name == "run_simulation", openWorldHint=False
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
_specs = [
|
||||
("run_simulation", "按水力步长在在线业务模型上执行指定时间段的实时模拟，不复制数据库；更新该时段的模型条件及实时基线，返回指定节点和管段结果；duration单位为秒，不创建方案或正式分析运行，结果不是实测", ScenarioInput, "simulation.run", "analysis", False, True),
|
||||
("run_contaminant_simulation", "在隔离模型中模拟假设污染源的浓度传播；返回指定节点/管段结果，不反演污染源，不代表实测污染范围", ContaminantSimulationInput, "simulation.run", "analysis", False, True),
|
||||
("analyze_valve_closure", "在隔离模型中模拟关阀的水力影响，返回指定节点/管段结果；不操作真实阀门", ValveClosureInput, "simulation.run", "analysis", False, True),
|
||||
("analyze_flushing", "在隔离模型中评估指定排放节点、冲洗流量和关阀组合的水力结果；不是现场冲洗或消毒合格证明", FlushingInput, "simulation.run", "analysis", False, True),
|
||||
("analyze_pressure_regulation", "在隔离模型中评估给定调泵序列和水箱初始水位的压力影响；不自动优化调度，不实际控制泵", PressureRegulationInput, "simulation.run", "analysis", False, True),
|
||||
("optimize_sensor_placement", "使用灵敏度或K-means算法推荐压力监测点；仅返回候选节点，不保存正式方案，不保证污染溯源可观测性", SensorPlacementInput, "optimization.run", "analysis", False, True),
|
||||
("run_leakage_identification", "基于SCADA压力与模型进行DMA漏损分配识别，需有依据的总漏损流量和时间对齐；返回区域估计，不是已确认漏点。内部执行遗传算法与水力计算", LeakageIdentificationInput, "burst.run", "analysis", False, True),
|
||||
("run_burst_detection", "爆管异常侦测：使用至少5个完整压力传感器此前15天数据；只提示压力异常，不定位、不证明爆管", BurstDetectionInput, "burst.run", "analysis", False, True),
|
||||
("run_burst_location", "爆管候选定位：对比异常与正常压力，可辅助流量；需有依据的漏水量假设。压力节点可省略，自动读取项目SCADA压力测点；monitoring正常窗可成对省略，默认前一天同时间段，须说明该对照假设。simulation须有已有分析运行UUID并读取同窗模型基线，不是实测；定位内部会执行水力计算，不创建新的事件方案", BurstLocationInput, "burst.run", "analysis", False, True),
|
||||
("analyze_valve_isolation", "按事故候选要素分析拓扑隔离边界、建议关闭阀门和受影响节点；不是污染扩散或水力影响范围，不执行关阀，须现场核实", ValveIsolationInput, "simulation.run", "analysis", False, True),
|
||||
(
|
||||
"get_element_properties",
|
||||
"读取当前项目的单个管网元素属性",
|
||||
m.ElementInput,
|
||||
"webgis.view",
|
||||
"asset",
|
||||
),
|
||||
(
|
||||
"list_elements",
|
||||
"分页读取管网元素，不自动读取全网",
|
||||
m.ElementPageInput,
|
||||
"webgis.view",
|
||||
"asset",
|
||||
),
|
||||
(
|
||||
"query_nearby_elements",
|
||||
"查询坐标附近的模型元素，默认pipe；center为EPSG:4326 [经度,纬度]，可直接使用geocode.center；radius_m默认200米、最大10000米，limit默认5、最大100，支持offset。按点到完整几何的WGS84椭球最短距离筛选并排序，返回element_id、distance_m和属性；不查询原始GIS管线、不自动扩大半径、不解析道路交叉点",
|
||||
m.NearbyElementsInput,
|
||||
"webgis.view",
|
||||
"asset",
|
||||
),
|
||||
(
|
||||
"get_network_options",
|
||||
"读取模型时间、能源或网络选项",
|
||||
m.OptionsInput,
|
||||
"webgis.view",
|
||||
"asset",
|
||||
),
|
||||
(
|
||||
"get_network_schema",
|
||||
"读取当前支持的元素、SCADA和选项结构",
|
||||
m.SchemaInput,
|
||||
"webgis.view",
|
||||
"asset",
|
||||
),
|
||||
(
|
||||
"list_scada_devices",
|
||||
"分页读取SCADA设备和物理单位",
|
||||
m.PageInput,
|
||||
"scada.view",
|
||||
"asset",
|
||||
),
|
||||
("get_scada_device", "读取一个SCADA设备", m.DeviceInput, "scada.view", "asset"),
|
||||
(
|
||||
"query_scada_readings",
|
||||
"查询带时区且最多24小时的原始或清洗监测数据",
|
||||
m.ScadaInput,
|
||||
"scada.view",
|
||||
"observed",
|
||||
),
|
||||
(
|
||||
"query_model_baseline_series",
|
||||
"读取最多20个元素的模型基线，不能称为实测",
|
||||
m.BaselineSeriesInput,
|
||||
"simulation.view",
|
||||
"model_baseline",
|
||||
),
|
||||
(
|
||||
"get_model_baseline_at_time",
|
||||
"读取单个元素指定时刻的模型基线",
|
||||
m.BaselineAtTimeInput,
|
||||
"simulation.view",
|
||||
"model_baseline",
|
||||
),
|
||||
(
|
||||
"query_model_baseline_field",
|
||||
"分页读取指定时刻的模型字段",
|
||||
m.BaselineFieldInput,
|
||||
"simulation.view",
|
||||
"model_baseline",
|
||||
),
|
||||
(
|
||||
"query_analysis_series",
|
||||
"读取指定分析运行和元素字段时序",
|
||||
m.AnalysisSeriesInput,
|
||||
"simulation.view",
|
||||
"analysis",
|
||||
),
|
||||
(
|
||||
"query_analysis_values",
|
||||
"分页读取指定分析运行某时刻的元素字段",
|
||||
m.AnalysisValuesInput,
|
||||
"simulation.view",
|
||||
"analysis",
|
||||
),
|
||||
(
|
||||
"list_analysis_runs",
|
||||
"分页读取持久分析运行",
|
||||
m.RunsInput,
|
||||
"webgis.view",
|
||||
"analysis",
|
||||
),
|
||||
("get_analysis_run", "读取单个持久分析运行", m.RunInput, "webgis.view", "analysis"),
|
||||
(
|
||||
"list_analysis_results",
|
||||
"分页读取分析成果，可筛选result_type",
|
||||
m.ResultsInput,
|
||||
"webgis.view",
|
||||
"analysis",
|
||||
),
|
||||
(
|
||||
"list_sensor_placement_runs",
|
||||
"分页读取已有监测点选址运行",
|
||||
m.PageInput,
|
||||
"optimization.run",
|
||||
"analysis",
|
||||
),
|
||||
(
|
||||
"get_sensor_placement_run",
|
||||
"读取已有监测点选址成果",
|
||||
m.RunInput,
|
||||
"optimization.run",
|
||||
"analysis",
|
||||
),
|
||||
(
|
||||
"query_element_history",
|
||||
"比较实测与模型基线或分析结果；analysis_comparison要求run_id；limit/offset按序列、时间排序的数据点分页",
|
||||
m.HistoryInput,
|
||||
"scada.view",
|
||||
"observed",
|
||||
),
|
||||
(
|
||||
"get_element_gis_properties",
|
||||
"读取元素关联GIS台账",
|
||||
m.ElementInput,
|
||||
"webgis.view",
|
||||
"asset",
|
||||
),
|
||||
(
|
||||
"get_source_gis_properties",
|
||||
"读取指定图层源要素台账",
|
||||
m.SourceGisInput,
|
||||
"webgis.view",
|
||||
"asset",
|
||||
),
|
||||
(
|
||||
"get_workbench_runtime",
|
||||
"读取工作台已有运行态",
|
||||
m.StrictInput,
|
||||
"simulation.view",
|
||||
"task",
|
||||
),
|
||||
(
|
||||
"get_task_context",
|
||||
"读取绑定任务的有界摘要及handling_skill元数据；section=skill读取Server按任务workflow分发的业务处置Skill正文与版本，属于流程指引而非事件证据；其他section选择task/report/evidence/actions/progress/access_requests/plan/run，data_path按返回路径读取嵌套字段，limit/offset分页；report_revision保留选定报告版本",
|
||||
m.TaskContextInput,
|
||||
"operational_task.view",
|
||||
"task",
|
||||
True,
|
||||
),
|
||||
(
|
||||
"query_peripheral_data",
|
||||
"查询有限外围数据源；synthetic是合成证据",
|
||||
m.PeripheralInput,
|
||||
"simulation.view",
|
||||
"peripheral",
|
||||
),
|
||||
(
|
||||
"get_task_evidence",
|
||||
"读取当前绑定任务已保存的证据快照",
|
||||
m.EvidenceInput,
|
||||
"operational_task.view",
|
||||
"task",
|
||||
True,
|
||||
),
|
||||
]
|
||||
TOOLS = {row[0]: QueryTool(*row) for row in _specs}
|
||||
@@ -0,0 +1,76 @@
|
||||
"""MCP is opt-in; an incomplete OAuth resource must never start listening."""
|
||||
|
||||
from pathlib import Path
|
||||
from ipaddress import IPv4Address, IPv4Network
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from pydantic import Field, model_validator
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class McpSettings(BaseSettings):
|
||||
enabled: bool = False
|
||||
issuer_uri: str = ""
|
||||
dev_http_issuer_uri: str = ""
|
||||
environment: str = Field(default="production", validation_alias="ENVIRONMENT")
|
||||
resource_uri: str = ""
|
||||
audience: str = "tjwater-mcp"
|
||||
host_client_id: str = "tjwater-agent-mcp"
|
||||
allowed_hosts: list[str] = Field(
|
||||
default_factory=lambda: ["127.0.0.1:*", "localhost:*"]
|
||||
)
|
||||
allowed_origins: list[str] = Field(default_factory=list)
|
||||
context_ttl_seconds: int = Field(default=600, ge=30, le=600)
|
||||
query_timeout_seconds: int = Field(default=30, ge=1, le=60)
|
||||
max_calculation_timeout_seconds: int = Field(default=3600, ge=120, le=86400)
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix="TJWATER_MCP_",
|
||||
extra="ignore",
|
||||
populate_by_name=True,
|
||||
env_file=Path(__file__).resolve().parents[2] / ".env",
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_resource(self):
|
||||
if self.dev_http_issuer_uri:
|
||||
dev_url = urlsplit(self.dev_http_issuer_uri)
|
||||
try:
|
||||
address = IPv4Address(dev_url.hostname or "")
|
||||
private_ipv4 = any(address in IPv4Network(network) for network in (
|
||||
"10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16"
|
||||
))
|
||||
except ValueError:
|
||||
private_ipv4 = False
|
||||
if (
|
||||
self.environment != "development"
|
||||
or dev_url.scheme != "http"
|
||||
or not private_ipv4
|
||||
or dev_url.username or dev_url.password or dev_url.query or dev_url.fragment
|
||||
or self.dev_http_issuer_uri != self.issuer_uri
|
||||
):
|
||||
raise ValueError("development HTTP issuer requires development environment and an exact private IPv4 issuer")
|
||||
if not self.enabled:
|
||||
return self
|
||||
for index, value in enumerate((self.issuer_uri, self.resource_uri)):
|
||||
url = urlsplit(value)
|
||||
if (
|
||||
url.scheme not in {"http", "https"}
|
||||
or not url.hostname
|
||||
or url.username
|
||||
or url.password
|
||||
or url.query
|
||||
or url.fragment
|
||||
):
|
||||
raise ValueError("MCP requires explicit issuer and resource URIs")
|
||||
if url.scheme == "http" and url.hostname not in {
|
||||
"localhost",
|
||||
"127.0.0.1",
|
||||
"::1",
|
||||
}:
|
||||
if index != 0 or value != self.dev_http_issuer_uri:
|
||||
raise ValueError("non-loopback MCP and OAuth require HTTPS")
|
||||
if urlsplit(self.resource_uri).path != "/mcp":
|
||||
raise ValueError("MCP resource URI must identify the /mcp endpoint")
|
||||
if not self.audience.strip() or not self.host_client_id.strip():
|
||||
raise ValueError("MCP audience and host client ID are required")
|
||||
return self
|
||||
@@ -0,0 +1,112 @@
|
||||
"""An OAuth identity and a Server-owned binding are both required for data queries."""
|
||||
|
||||
import secrets
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from uuid import UUID
|
||||
from fastapi import HTTPException
|
||||
from app.auth.permissions import (
|
||||
permissions_for_context,
|
||||
OPERATIONAL_TASK_EXECUTE,
|
||||
OPERATIONAL_TASK_VIEW,
|
||||
)
|
||||
from app.infra.db.postgresql.operational_tasks import OperationalTaskRepository
|
||||
from app.infra.db.postgresql.mcp_queries import McpQueryRepository
|
||||
from app.mcp.models import BindInput, BoundContext
|
||||
from app.mcp.scope import QueryScope
|
||||
from app.mcp.config import McpSettings
|
||||
from mcp.server.auth.provider import AccessToken
|
||||
|
||||
|
||||
def task_repository(scope: QueryScope) -> OperationalTaskRepository:
|
||||
return OperationalTaskRepository(scope.connection, scope.project.project_id)
|
||||
|
||||
|
||||
async def verify_task(scope: QueryScope, binding: BoundContext) -> None:
|
||||
permissions = permissions_for_context(scope.project)
|
||||
if OPERATIONAL_TASK_VIEW not in permissions:
|
||||
raise HTTPException(403, "Task read permission is required")
|
||||
repo = task_repository(scope)
|
||||
async with scope.connection.transaction():
|
||||
task = await repo.get_task(str(binding.task_id))
|
||||
if not task:
|
||||
raise HTTPException(404, "Task is not in this project")
|
||||
if binding.purpose == "task_collaboration":
|
||||
return
|
||||
if OPERATIONAL_TASK_EXECUTE not in permissions:
|
||||
raise HTTPException(403, "Task execution permission is required")
|
||||
run = await repo.get_active_run(str(binding.task_id))
|
||||
now = datetime.now(timezone.utc)
|
||||
owner = binding.parent_session_id or binding.native_session_id
|
||||
if (
|
||||
not run
|
||||
or run["status"] not in {"claimed", "running"}
|
||||
or run.get("agent_session_id") != owner
|
||||
or not run.get("lease_expires_at")
|
||||
or run["lease_expires_at"] <= now
|
||||
or binding.run_id is not None
|
||||
and str(run["run_id"]) != str(binding.run_id)
|
||||
or binding.execution_request_id is not None
|
||||
and str(run["execution_request_id"]) != str(binding.execution_request_id)
|
||||
):
|
||||
raise HTTPException(
|
||||
409, "Investigation run is no longer active for this session"
|
||||
)
|
||||
binding.run_id = UUID(str(run["run_id"]))
|
||||
binding.execution_request_id = UUID(str(run["execution_request_id"]))
|
||||
|
||||
|
||||
async def bind_context(
|
||||
scope: QueryScope, principal: AccessToken, command: BindInput, settings: McpSettings
|
||||
) -> BoundContext:
|
||||
if principal.client_id != settings.host_client_id:
|
||||
raise HTTPException(
|
||||
403, "Only the trusted product host may bind query contexts"
|
||||
)
|
||||
from app.mcp.catalog import TOOLS
|
||||
|
||||
binding = BoundContext(
|
||||
context_handle=secrets.token_urlsafe(32),
|
||||
project_id=scope.project.project_id,
|
||||
subject=principal.subject,
|
||||
client_id=principal.client_id,
|
||||
tool_cap=sorted(name for name, tool in TOOLS.items() if command.purpose != "task_reader" or not tool.computation),
|
||||
permission_cap=sorted(permissions_for_context(scope.project)),
|
||||
native_session_id=command.native_session_id,
|
||||
purpose=command.purpose,
|
||||
parent_session_id=command.parent_session_id,
|
||||
task_id=command.task_id,
|
||||
expires_at=datetime.now(timezone.utc)
|
||||
+ timedelta(seconds=settings.context_ttl_seconds),
|
||||
)
|
||||
if binding.task_id:
|
||||
await verify_task(scope, binding)
|
||||
from app.infra.db.postgresql.read_queries import task_data_scopes
|
||||
|
||||
async with scope.connection.transaction():
|
||||
binding.task_scope_cap = await task_data_scopes(scope.connection,
|
||||
project_id=binding.project_id, task_id=binding.task_id,
|
||||
permissions=permissions_for_context(scope.project))
|
||||
await McpQueryRepository(scope.connection, scope.project.project_id).create(binding)
|
||||
return binding
|
||||
|
||||
|
||||
async def resolve_context(
|
||||
scope: QueryScope, principal: AccessToken, handle: str
|
||||
) -> BoundContext:
|
||||
binding = await McpQueryRepository(scope.connection, scope.project.project_id).get(
|
||||
handle
|
||||
)
|
||||
if (
|
||||
not binding
|
||||
or binding.revoked_at
|
||||
or binding.expires_at <= datetime.now(timezone.utc)
|
||||
or binding.subject != principal.subject
|
||||
or binding.client_id != principal.client_id
|
||||
or binding.project_id != scope.project.project_id
|
||||
):
|
||||
raise HTTPException(
|
||||
403, "Query context is unavailable for this identity and project"
|
||||
)
|
||||
if binding.task_id:
|
||||
await verify_task(scope, binding)
|
||||
return binding
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Delegate MCP transport/metadata with the original request path intact."""
|
||||
|
||||
|
||||
class McpDispatchMiddleware:
|
||||
def __init__(self, app, runtime):
|
||||
self.app, self.runtime = app, runtime
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
path = scope.get("path", "")
|
||||
if scope["type"] == "http" and (
|
||||
path == "/mcp" or path.startswith("/.well-known/oauth-protected-resource")
|
||||
):
|
||||
await self.runtime.asgi(scope, receive, send)
|
||||
else:
|
||||
await self.app(scope, receive, send)
|
||||
@@ -0,0 +1,293 @@
|
||||
"""Strict business inputs and provenance envelopes shared by MCP handlers."""
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Annotated, Literal
|
||||
from uuid import UUID
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator
|
||||
|
||||
Id = Annotated[str, Field(min_length=1, max_length=160, pattern=r"^\S(?:.*\S)?$")]
|
||||
ElementType = Literal["junction", "pipe", "valve", "pump", "tank", "reservoir"]
|
||||
SeriesType = Literal["junction", "pipe"]
|
||||
|
||||
|
||||
class StrictInput(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
context_handle: Annotated[str, Field(min_length=32, max_length=128)]
|
||||
|
||||
|
||||
class PageInput(StrictInput):
|
||||
limit: int = Field(default=20, ge=1, le=500)
|
||||
offset: int = Field(default=0, ge=0, le=10000)
|
||||
|
||||
|
||||
class TaskContextInput(PageInput):
|
||||
section: Literal["summary", "task", "report", "evidence", "actions", "progress", "access_requests", "plan", "run", "skill"] = "summary"
|
||||
report_revision: int | None = Field(default=None, ge=1)
|
||||
data_path: list[Id] = Field(default_factory=list, max_length=8)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def summary_is_fixed(self):
|
||||
if self.section in {"summary", "skill"} and (self.data_path or self.offset):
|
||||
raise ValueError("Select a detail section before using data_path or offset")
|
||||
return self
|
||||
|
||||
|
||||
class ElementInput(StrictInput):
|
||||
element_type: ElementType
|
||||
element_id: Id
|
||||
|
||||
|
||||
class ElementPageInput(PageInput):
|
||||
element_type: ElementType
|
||||
|
||||
|
||||
class NearbyElementsInput(PageInput):
|
||||
element_type: ElementType = "pipe"
|
||||
center: tuple[
|
||||
Annotated[float, Field(ge=-180, le=180, allow_inf_nan=False)],
|
||||
Annotated[float, Field(ge=-90, le=90, allow_inf_nan=False)],
|
||||
] = Field(description="EPSG:4326 [longitude, latitude], directly from geocode.center")
|
||||
radius_m: float = Field(default=200, gt=0, le=10000, allow_inf_nan=False)
|
||||
limit: int = Field(default=5, ge=1, le=100)
|
||||
|
||||
|
||||
class OptionsInput(StrictInput):
|
||||
kind: Literal["time", "energy", "pump-energy", "network"]
|
||||
pump_id: Id | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def pump_required(self):
|
||||
if self.kind == "pump-energy" and not self.pump_id:
|
||||
raise ValueError("pump_id is required for pump-energy")
|
||||
if self.kind != "pump-energy" and self.pump_id:
|
||||
raise ValueError("pump_id is only valid for pump-energy")
|
||||
return self
|
||||
|
||||
|
||||
class SchemaInput(StrictInput):
|
||||
kind: Literal[
|
||||
"junction",
|
||||
"pipe",
|
||||
"valve",
|
||||
"pump",
|
||||
"tank",
|
||||
"reservoir",
|
||||
"time",
|
||||
"energy",
|
||||
"pump-energy",
|
||||
"option",
|
||||
"scada-device",
|
||||
]
|
||||
|
||||
|
||||
class DeviceInput(StrictInput):
|
||||
device_id: Id
|
||||
|
||||
|
||||
class WindowInput(PageInput):
|
||||
start: datetime
|
||||
end: datetime
|
||||
limit: int = Field(default=20, ge=1, le=10000)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def window_valid(self):
|
||||
if self.start.tzinfo is None or self.end.tzinfo is None:
|
||||
raise ValueError("timestamps require explicit timezone")
|
||||
if not timedelta(0) < self.end - self.start <= timedelta(days=1):
|
||||
raise ValueError("time window must be positive and at most 24 hours")
|
||||
return self
|
||||
|
||||
|
||||
class ScadaInput(WindowInput):
|
||||
device_ids: list[Id] = Field(min_length=1, max_length=20)
|
||||
field: Literal["monitored_value", "cleaned_value"] = "monitored_value"
|
||||
|
||||
|
||||
class BaselineSeriesInput(WindowInput):
|
||||
element_type: SeriesType
|
||||
element_ids: list[Id] = Field(min_length=1, max_length=20)
|
||||
|
||||
|
||||
class TimeInput(StrictInput):
|
||||
time: datetime
|
||||
|
||||
@model_validator(mode="after")
|
||||
def timezone_required(self):
|
||||
if self.time.tzinfo is None:
|
||||
raise ValueError("timestamp requires explicit timezone")
|
||||
return self
|
||||
|
||||
|
||||
class BaselineAtTimeInput(TimeInput):
|
||||
element_type: SeriesType
|
||||
element_id: Id
|
||||
|
||||
|
||||
class BaselineFieldInput(TimeInput):
|
||||
element_type: SeriesType
|
||||
field: Id
|
||||
limit: int = Field(default=20, ge=1, le=500)
|
||||
offset: int = Field(default=0, ge=0, le=10000)
|
||||
|
||||
|
||||
class AnalysisSeriesInput(WindowInput):
|
||||
run_id: UUID
|
||||
element_type: SeriesType
|
||||
element_id: Id
|
||||
field: Id
|
||||
|
||||
|
||||
class AnalysisValuesInput(BaselineFieldInput):
|
||||
run_id: UUID
|
||||
|
||||
|
||||
class RunInput(StrictInput):
|
||||
run_id: UUID
|
||||
|
||||
|
||||
class RunsInput(PageInput):
|
||||
limit: int = Field(default=20, ge=1, le=100)
|
||||
|
||||
|
||||
class ResultsInput(PageInput):
|
||||
run_id: UUID
|
||||
result_type: Id | None = None
|
||||
|
||||
|
||||
class HistoryTarget(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
element_type: SeriesType
|
||||
element_id: Id
|
||||
device_ids: list[Id] | None = Field(default=None, min_length=1, max_length=20)
|
||||
|
||||
|
||||
class HistoryInput(WindowInput):
|
||||
elements: list[HistoryTarget] = Field(min_length=1, max_length=20)
|
||||
mode: Literal["observed", "realtime_comparison", "analysis_comparison"] = "observed"
|
||||
run_id: UUID | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def analysis_required(self):
|
||||
if (self.mode == "analysis_comparison") != (self.run_id is not None):
|
||||
raise ValueError("run_id is required only for analysis_comparison")
|
||||
return self
|
||||
|
||||
|
||||
class SourceGisInput(StrictInput):
|
||||
layer_key: Annotated[str, Field(min_length=1, max_length=80)]
|
||||
source_fid: int = Field(ge=0, le=9223372036854775807)
|
||||
|
||||
|
||||
class PeripheralInput(PageInput):
|
||||
source: Literal[
|
||||
"peripheral_points",
|
||||
"peripheral_work_orders",
|
||||
"peripheral_observations",
|
||||
"peripheral_test_data",
|
||||
]
|
||||
point_id: Id | None = None
|
||||
start: datetime | None = None
|
||||
end: datetime | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def observation_window(self):
|
||||
if self.source == "peripheral_observations":
|
||||
if not self.point_id or self.start is None or self.end is None:
|
||||
raise ValueError("observations require point_id, start and end")
|
||||
if (
|
||||
self.start.tzinfo is None
|
||||
or self.end.tzinfo is None
|
||||
or not timedelta(0) < self.end - self.start <= timedelta(days=1)
|
||||
):
|
||||
raise ValueError(
|
||||
"observation window requires timezone and at most 24 hours"
|
||||
)
|
||||
elif (
|
||||
self.point_id is not None or self.start is not None or self.end is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"point_id and window are only valid for peripheral_observations"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class EvidenceInput(PageInput):
|
||||
evidence_id: Id
|
||||
data_path: list[Id] = Field(default_factory=list, max_length=8)
|
||||
|
||||
|
||||
class BindInput(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
native_session_id: Id
|
||||
purpose: Literal[
|
||||
"ordinary_chat", "task_execution", "task_reader", "task_collaboration"
|
||||
]
|
||||
task_id: UUID | None = None
|
||||
parent_session_id: Id | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def task_scope(self):
|
||||
if self.purpose == "ordinary_chat" and (self.task_id or self.parent_session_id):
|
||||
raise ValueError("ordinary chat cannot bind a task")
|
||||
if self.purpose != "ordinary_chat" and self.task_id is None:
|
||||
raise ValueError("task context requires task_id")
|
||||
if (self.purpose == "task_reader") != (self.parent_session_id is not None):
|
||||
raise ValueError("only a task reader requires parent_session_id")
|
||||
return self
|
||||
|
||||
|
||||
class BindResponse(BaseModel):
|
||||
context_handle: str
|
||||
expires_at: datetime
|
||||
task_id: UUID | None
|
||||
run_id: UUID | None
|
||||
|
||||
|
||||
class PageInfo(BaseModel):
|
||||
limit: int
|
||||
offset: int
|
||||
total: int | None = None
|
||||
has_more: bool
|
||||
|
||||
|
||||
class QueryResult(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
source: str
|
||||
source_ref: str
|
||||
queried_at: datetime
|
||||
origin: Literal["system_data", "synthetic"] = "system_data"
|
||||
semantics: Literal[
|
||||
"asset",
|
||||
"observed",
|
||||
"model_baseline",
|
||||
"analysis",
|
||||
"comparison",
|
||||
"task",
|
||||
"peripheral",
|
||||
]
|
||||
data: JsonValue
|
||||
page: PageInfo | None = None
|
||||
window: dict[str, datetime] | None = None
|
||||
units: dict[str, str] = Field(default_factory=dict)
|
||||
truncated: bool = False
|
||||
data_gaps: list[str] = Field(default_factory=list)
|
||||
evidence_id: str | None = None
|
||||
|
||||
|
||||
class BoundContext(BaseModel):
|
||||
context_handle: str
|
||||
project_id: UUID
|
||||
subject: str
|
||||
client_id: str
|
||||
tool_cap: list[str] = Field(default_factory=list)
|
||||
permission_cap: list[str] = Field(default_factory=list)
|
||||
task_scope_cap: list[str] = Field(default_factory=list)
|
||||
native_session_id: str
|
||||
purpose: str
|
||||
parent_session_id: str | None = None
|
||||
task_id: UUID | None = None
|
||||
run_id: UUID | None = None
|
||||
execution_request_id: UUID | None = None
|
||||
expires_at: datetime
|
||||
revoked_at: datetime | None = None
|
||||
@@ -0,0 +1,132 @@
|
||||
"""Read existing project-owned peripheral facts without forwarding MCP tokens."""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from uuid import uuid4
|
||||
from fastapi import HTTPException
|
||||
from app.infra.db.postgresql.read_queries import MAX_QUERY_BYTES, read_rows, scada_devices, series_rows
|
||||
from app.domain.schemas.operational_data import TaskDataSnapshot, TaskDataQuery
|
||||
from app.services.content_digest import digest
|
||||
from samples.peripheral.store import PostgresStore
|
||||
from app.infra.integrations.operational_data import _select_demo_observations
|
||||
|
||||
|
||||
async def read_peripheral_query(scope, args):
|
||||
conn = scope.connection
|
||||
project = str(scope.project.project_id)
|
||||
truncated = False
|
||||
origin = "system_data"
|
||||
if args.source in {"peripheral_work_orders", "peripheral_test_data"}:
|
||||
table, column = (
|
||||
("work_orders", "work_order_id")
|
||||
if args.source == "peripheral_work_orders"
|
||||
else ("test_data_records", "record_id")
|
||||
)
|
||||
rows = await read_rows(
|
||||
conn,
|
||||
"SELECT document FROM peripheral_sample."
|
||||
+ table
|
||||
+ " WHERE project_id=%s AND deleted_at IS NULL ORDER BY "
|
||||
+ column
|
||||
+ " LIMIT %s OFFSET %s",
|
||||
(project, args.limit + 1, args.offset),
|
||||
)
|
||||
items = [
|
||||
PostgresStore._active_managed_document(row["document"])
|
||||
for row in rows[: args.limit]
|
||||
]
|
||||
truncated = len(rows) > args.limit
|
||||
if any(item.get("project_id") != project for item in items):
|
||||
raise HTTPException(502, "Peripheral project mismatch")
|
||||
data = {
|
||||
"items": items,
|
||||
"truncated": truncated,
|
||||
"limit": args.limit,
|
||||
"offset": args.offset,
|
||||
}
|
||||
origin = (
|
||||
"synthetic"
|
||||
if items and any(item.get("origin") == "synthetic" for item in items)
|
||||
else "system_data"
|
||||
)
|
||||
elif args.source == "peripheral_points":
|
||||
rows = await scada_devices(conn, limit=args.limit + 1, offset=args.offset)
|
||||
data = {
|
||||
"origin": "historical_replay",
|
||||
"items": [
|
||||
{
|
||||
"point_id": row["device_id"],
|
||||
"label": row["device_id"] + " · " + row["device_type"],
|
||||
"unit": row["measurement_unit"],
|
||||
"status": "unknown",
|
||||
"origin": "historical_replay",
|
||||
}
|
||||
for row in rows[: args.limit]
|
||||
],
|
||||
"truncated": len(rows) > args.limit,
|
||||
}
|
||||
truncated = data["truncated"]
|
||||
elif args.point_id.startswith("DEMO-"):
|
||||
rows = await read_rows(
|
||||
conn,
|
||||
"SELECT CASE WHEN octet_length(document::text)<=%s THEN document END AS document,"
|
||||
"digest,octet_length(document::text)>%s AS too_large "
|
||||
"FROM peripheral_sample.demo_datasets WHERE project_id=%s AND dataset_id='water-demo-v1'",
|
||||
(MAX_QUERY_BYTES, MAX_QUERY_BYTES, project),
|
||||
)
|
||||
if rows and rows[0]["too_large"]:
|
||||
raise HTTPException(413, "Synthetic dataset exceeds the integrity-check byte budget")
|
||||
if not rows or digest(rows[0]["document"]) != rows[0]["digest"]:
|
||||
raise HTTPException(503, "Synthetic dataset is unavailable or invalid")
|
||||
query = TaskDataQuery(
|
||||
source=args.source, point_id=args.point_id, start=args.start, end=args.end
|
||||
)
|
||||
data = _select_demo_observations(
|
||||
{**rows[0]["document"], "project_id": project}, query, project
|
||||
)
|
||||
items = data["items"]
|
||||
truncated = len(items) > args.limit + args.offset
|
||||
data["items"] = items[args.offset : args.offset + args.limit]
|
||||
data["truncated"] = truncated
|
||||
origin = "synthetic"
|
||||
else:
|
||||
devices = await scada_devices(conn, device_id=args.point_id, limit=1)
|
||||
if not devices:
|
||||
raise HTTPException(404, "Peripheral point was not found")
|
||||
device = devices[0]
|
||||
if not device.get("api_query_id"):
|
||||
raise HTTPException(503, "Peripheral history mapping is unavailable")
|
||||
async with scope.timescale() as ts:
|
||||
rows = await series_rows(
|
||||
ts,
|
||||
source="scada",
|
||||
element_type="junction",
|
||||
ids=[device["api_query_id"]],
|
||||
start=args.start,
|
||||
end=args.end,
|
||||
field="monitored_value",
|
||||
limit=args.limit + 1,
|
||||
offset=args.offset,
|
||||
end_exclusive=True,
|
||||
)
|
||||
truncated = len(rows) > args.limit
|
||||
data = {
|
||||
"origin": "historical_replay",
|
||||
"items": [
|
||||
{
|
||||
"point_id": args.point_id,
|
||||
"observed_at": row["time"].isoformat(),
|
||||
"value": row["value"],
|
||||
"unit": device["measurement_unit"],
|
||||
}
|
||||
for row in rows[: args.limit]
|
||||
],
|
||||
"truncated": truncated,
|
||||
}
|
||||
return TaskDataSnapshot(
|
||||
source=args.source,
|
||||
source_ref="mcp:" + str(uuid4()),
|
||||
queried_at=datetime.now(timezone.utc),
|
||||
data=data,
|
||||
origin=origin,
|
||||
truncated=truncated,
|
||||
)
|
||||
@@ -0,0 +1,492 @@
|
||||
"""MCP domain dispatch reuses Server repositories and services, never an HTTP loopback."""
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
from contextlib import nullcontext
|
||||
from typing import Any
|
||||
from pydantic import BaseModel
|
||||
from app.mcp.scope import QueryScope
|
||||
from app.mcp.catalog import QueryTool, TOOLS
|
||||
from app.mcp.models import BoundContext
|
||||
from datetime import datetime, timezone
|
||||
from uuid import uuid4
|
||||
from fastapi import HTTPException
|
||||
from fastapi.encoders import jsonable_encoder
|
||||
from app.auth.permissions import permissions_for_context
|
||||
from app.infra.db.postgresql import read_queries as reads
|
||||
from app.infra.db.postgresql.spatial_queries import nearby_elements
|
||||
from app.infra.db.postgresql.analysis import AnalysisRepository
|
||||
from app.infra.db.postgresql.gis_attributes import GisAttributeRepository
|
||||
from app.infra.db.postgresql.network_settings import NetworkSettingsRepository
|
||||
from app.infra.db.postgresql.operational_tasks import OperationalTaskRepository
|
||||
from app.services.operational_tasks import OperationalTaskService
|
||||
from app.mcp.models import QueryResult, PageInfo
|
||||
from app.mcp.scenario_service import SCENARIO_TOOLS
|
||||
from app.mcp.work_order_skills import read_work_order_skill
|
||||
|
||||
|
||||
def finite_data(value: Any) -> tuple[Any, int]:
|
||||
"""JSON cannot represent infinity/NaN; preserve the gap as null with provenance."""
|
||||
if isinstance(value, float) and not math.isfinite(value):
|
||||
return None, 1
|
||||
count = 0
|
||||
if isinstance(value, dict):
|
||||
result = {}
|
||||
for key, item in value.items():
|
||||
result[key], replaced = finite_data(item)
|
||||
count += replaced
|
||||
return result, count
|
||||
if isinstance(value, list):
|
||||
result = []
|
||||
for item in value:
|
||||
clean, replaced = finite_data(item)
|
||||
result.append(clean)
|
||||
count += replaced
|
||||
return result, count
|
||||
return value, 0
|
||||
|
||||
|
||||
def task_service(scope: QueryScope) -> OperationalTaskService:
|
||||
return OperationalTaskService(
|
||||
OperationalTaskRepository(scope.connection, scope.project.project_id),
|
||||
scope.project,
|
||||
set(permissions_for_context(scope.project)),
|
||||
)
|
||||
|
||||
|
||||
def page_result(rows, args):
|
||||
has_more = len(rows) > args.limit
|
||||
return (
|
||||
rows[: args.limit],
|
||||
PageInfo(limit=args.limit, offset=args.offset, total=None, has_more=has_more),
|
||||
has_more,
|
||||
)
|
||||
|
||||
|
||||
async def authorize_query(
|
||||
scope: QueryScope, binding: BoundContext, tool: QueryTool, args: BaseModel
|
||||
) -> None:
|
||||
if tool.name not in binding.tool_cap:
|
||||
raise HTTPException(403, "Query is outside the captured tool catalog")
|
||||
permissions = permissions_for_context(scope.project)
|
||||
if tool.computation:
|
||||
if tool.name == "optimize_sensor_placement" and scope.project.project_role != "member":
|
||||
raise HTTPException(403, "Sensor placement requires project member access")
|
||||
if binding.purpose == "task_reader":
|
||||
raise HTTPException(403, "Read-only task agents cannot start calculations")
|
||||
required = {tool.permission}
|
||||
if tool.name in SCENARIO_TOOLS:
|
||||
required.add("scada.view")
|
||||
if tool.name in {"run_burst_detection", "run_burst_location", "run_leakage_identification"}:
|
||||
required.add("simulation.view" if getattr(args, "data_source", "monitoring") == "simulation" else "scada.view")
|
||||
if not required.issubset(permissions) or not required.issubset(binding.permission_cap):
|
||||
raise HTTPException(403, "Calculation and input data permissions are required")
|
||||
permission = (
|
||||
"scada.view"
|
||||
if tool.name == "get_network_schema" and args.kind == "scada-device"
|
||||
else tool.permission
|
||||
)
|
||||
if tool.task_only and binding.task_id is None:
|
||||
raise HTTPException(403, "This query requires a bound task")
|
||||
granted = permission in permissions and permission in binding.permission_cap
|
||||
source = {
|
||||
"list_scada_devices": "scada_devices",
|
||||
"get_scada_device": "scada_devices",
|
||||
"query_scada_readings": "scada_readings",
|
||||
"query_element_history": "scada_readings",
|
||||
"get_workbench_runtime": "workbench_runtime",
|
||||
}.get(tool.name)
|
||||
if tool.name == "query_peripheral_data":
|
||||
source = args.source
|
||||
if not granted and binding.task_id and source in binding.task_scope_cap:
|
||||
async with scope.connection.transaction():
|
||||
granted = await task_service(scope).can_read_source(binding.task_id, source)
|
||||
if not granted:
|
||||
raise HTTPException(403, "Data source permission is required")
|
||||
if (
|
||||
tool.name == "query_element_history"
|
||||
and args.mode != "observed"
|
||||
and (
|
||||
"simulation.view" not in permissions
|
||||
or "simulation.view" not in binding.permission_cap
|
||||
)
|
||||
):
|
||||
raise HTTPException(403, "Model comparison permission is required")
|
||||
|
||||
|
||||
async def execute_query(
|
||||
scope: QueryScope,
|
||||
binding: BoundContext,
|
||||
tool: QueryTool,
|
||||
args: BaseModel,
|
||||
*,
|
||||
full_snapshot: bool = False,
|
||||
) -> QueryResult:
|
||||
await authorize_query(scope, binding, tool, args)
|
||||
name = tool.name
|
||||
conn = scope.connection
|
||||
page = None
|
||||
truncated = False
|
||||
units = {}
|
||||
origin = "system_data"
|
||||
gaps = []
|
||||
source_ref = "mcp:" + str(uuid4())
|
||||
evidence_id = None
|
||||
if tool.computation:
|
||||
from app.mcp.analysis_service import TIMESCALE_CALCULATIONS, calculate
|
||||
|
||||
# Context/receipt reads may have started an implicit transaction. Release
|
||||
# it before potentially long CPU/IO work; evidence uses a new transaction.
|
||||
await conn.commit()
|
||||
async with (
|
||||
scope.simulation_routing()
|
||||
if name in TIMESCALE_CALCULATIONS
|
||||
else nullcontext()
|
||||
):
|
||||
data = await calculate(
|
||||
name, args, network=scope.project.project_code,
|
||||
username=scope.user.username,
|
||||
)
|
||||
data = {
|
||||
"parameters": args.model_dump(mode="json", exclude={"context_handle"}),
|
||||
"result": data,
|
||||
}
|
||||
gaps = [{
|
||||
"run_simulation": "Online hydraulic recalculation updated the realtime baseline; results cover requested elements only, not observed SCADA data. No standalone analysis run was saved.",
|
||||
"run_contaminant_simulation": "Forward model simulation for an assumed source, not source inversion or observed contamination; only requested elements are returned.",
|
||||
"analyze_valve_closure": "Simulated hydraulic impact for requested elements only; no physical valves were operated.",
|
||||
"analyze_flushing": "Simulated flushing hydraulics only; no field action or water-quality clearance was verified.",
|
||||
"analyze_pressure_regulation": "Evaluation of a supplied pump schedule, not automatic optimization or device control; requested elements only.",
|
||||
"optimize_sensor_placement": "Candidate pressure monitoring locations, not a saved deployment scheme or a guarantee of contaminant source observability.",
|
||||
"run_leakage_identification": "DMA leakage estimates depend on the total leakage constraint and aligned model/SCADA time; they do not confirm an exact leak location.",
|
||||
"run_burst_detection": "Pressure anomalies alone do not confirm a burst or locate a pipe.",
|
||||
"run_burst_location": "Candidate ranking depends on the leakage assumption, model and sensor coverage; field confirmation is required.",
|
||||
"analyze_valve_isolation": "Topology-only isolation advice; valve operability, hydraulic impact and field actions are not verified or executed.",
|
||||
}[name]]
|
||||
if getattr(args, "data_source", None) == "simulation":
|
||||
gaps.append("Input is an existing analysis run and model baseline, not observed SCADA data.")
|
||||
elif name == "query_nearby_elements":
|
||||
rows = await nearby_elements(
|
||||
conn, args.element_type, center=args.center, radius_m=args.radius_m,
|
||||
limit=args.limit + 1, offset=args.offset,
|
||||
)
|
||||
items, page, truncated = page_result(rows, args)
|
||||
truncated = truncated or args.offset > 0
|
||||
data = {
|
||||
"source": "model_network", "element_type": args.element_type,
|
||||
"center": list(args.center), "center_crs": "EPSG:4326",
|
||||
"radius_m": args.radius_m, "distance_method": "WGS84_spheroid",
|
||||
"items": items,
|
||||
}
|
||||
units = {"distance_m": "m", "radius_m": "m"}
|
||||
gaps = ["Spatial proximity does not establish hydraulic connectivity or a supply relationship"]
|
||||
elif name in {"get_element_properties", "list_elements"}:
|
||||
rows = await reads.asset_rows(
|
||||
conn,
|
||||
args.element_type,
|
||||
element_id=getattr(args, "element_id", None),
|
||||
limit=getattr(args, "limit", 1) + 1,
|
||||
offset=getattr(args, "offset", 0),
|
||||
)
|
||||
if name == "get_element_properties":
|
||||
if not rows:
|
||||
raise HTTPException(404, "Element was not found")
|
||||
data = rows[0]
|
||||
else:
|
||||
data, page, truncated = page_result(rows, args)
|
||||
elif name in {"list_scada_devices", "get_scada_device"}:
|
||||
rows = await reads.scada_devices(
|
||||
conn,
|
||||
device_id=getattr(args, "device_id", None),
|
||||
limit=getattr(args, "limit", 1) + 1,
|
||||
offset=getattr(args, "offset", 0),
|
||||
)
|
||||
if name == "get_scada_device":
|
||||
if not rows:
|
||||
raise HTTPException(404, "SCADA device was not found")
|
||||
data = rows[0]
|
||||
else:
|
||||
data, page, truncated = page_result(rows, args)
|
||||
elif name == "get_network_schema":
|
||||
from app.services import tjnetwork
|
||||
|
||||
mapping = {
|
||||
"option": "get_option_v3_schema",
|
||||
"scada-device": "get_scada_info_schema",
|
||||
"pump-energy": "get_pump_energy_schema",
|
||||
}
|
||||
function = getattr(
|
||||
tjnetwork, mapping.get(args.kind, "get_" + args.kind + "_schema")
|
||||
)
|
||||
data = await asyncio.to_thread(function, scope.project.project_code)
|
||||
elif name == "get_network_options":
|
||||
if args.kind == "pump-energy":
|
||||
rows = await reads.read_rows(
|
||||
conn,
|
||||
"SELECT price,pattern_id AS pattern,efficiency_curve_id AS effic FROM network.pump_energy_settings WHERE pump_id=%s",
|
||||
(args.pump_id,),
|
||||
)
|
||||
data = {"pump": args.pump_id, "price": None, "pattern": None, "effic": None}
|
||||
if rows:
|
||||
data.update(rows[0])
|
||||
else:
|
||||
table = {
|
||||
"time": "time_settings",
|
||||
"energy": "energy_settings",
|
||||
"network": "simulation_settings",
|
||||
}[args.kind]
|
||||
where = " WHERE engine_version='v3'" if args.kind == "network" else ""
|
||||
rows = await reads.read_rows(
|
||||
conn,
|
||||
"SELECT key,value FROM network."
|
||||
+ table
|
||||
+ where
|
||||
+ " ORDER BY key LIMIT 501",
|
||||
)
|
||||
data = {str(row["key"]): str(row["value"]) for row in rows[:500]}
|
||||
truncated = len(rows) > 500
|
||||
elif name in {
|
||||
"query_scada_readings",
|
||||
"query_model_baseline_series",
|
||||
"get_model_baseline_at_time",
|
||||
"query_model_baseline_field",
|
||||
"query_analysis_series",
|
||||
"query_analysis_values",
|
||||
}:
|
||||
source = (
|
||||
"scada"
|
||||
if name == "query_scada_readings"
|
||||
else "analysis"
|
||||
if name.startswith("query_analysis")
|
||||
else "realtime"
|
||||
)
|
||||
ids = (
|
||||
getattr(args, "device_ids", None)
|
||||
or getattr(args, "element_ids", None)
|
||||
or ([args.element_id] if hasattr(args, "element_id") else [])
|
||||
)
|
||||
async with scope.timescale() as ts:
|
||||
rows = await reads.series_rows(
|
||||
ts,
|
||||
source=source,
|
||||
element_type=getattr(args, "element_type", "junction"),
|
||||
ids=ids,
|
||||
start=getattr(args, "start", None),
|
||||
end=getattr(args, "end", None),
|
||||
time=getattr(args, "time", None),
|
||||
field=getattr(args, "field", None),
|
||||
run_id=getattr(args, "run_id", None),
|
||||
limit=getattr(args, "limit", 1) + 1,
|
||||
offset=getattr(args, "offset", 0),
|
||||
)
|
||||
if name == "get_model_baseline_at_time":
|
||||
data = rows[:1]
|
||||
else:
|
||||
data, page, truncated = page_result(rows, args)
|
||||
if source != "scada":
|
||||
async with conn.transaction():
|
||||
units = await NetworkSettingsRepository.get_result_units(conn)
|
||||
if source == "analysis":
|
||||
run = await AnalysisRepository.get_run(conn, args.run_id)
|
||||
if not run:
|
||||
raise HTTPException(404, "Analysis run was not found")
|
||||
recorded = (run.get("parameters") or {}).get("result_units")
|
||||
if isinstance(recorded, dict):
|
||||
units = recorded
|
||||
else:
|
||||
gaps.append(
|
||||
"Analysis units inferred from current network settings"
|
||||
)
|
||||
if hasattr(args, "field") and args.field not in units:
|
||||
gaps.append(f"Unit metadata is unavailable for field {args.field}")
|
||||
else:
|
||||
metadata = [
|
||||
r
|
||||
for device in ids
|
||||
for r in await reads.scada_devices(conn, device_id=device, limit=1)
|
||||
]
|
||||
units = {
|
||||
r["device_id"]: r["measurement_unit"]
|
||||
for r in metadata
|
||||
if r["device_id"] in ids
|
||||
}
|
||||
elif name in {"list_analysis_runs", "get_analysis_run", "list_analysis_results"}:
|
||||
async with conn.transaction():
|
||||
if name == "get_analysis_run":
|
||||
data = await AnalysisRepository.get_run(conn, args.run_id)
|
||||
if not data:
|
||||
raise HTTPException(404, "Analysis run was not found")
|
||||
elif name == "list_analysis_runs":
|
||||
# Fetch one additional row with the repository's supported summary-page bound.
|
||||
rows = await AnalysisRepository.list_run_summaries(
|
||||
conn, limit=args.limit + 1, offset=args.offset
|
||||
)
|
||||
data, page, truncated = page_result(rows, args)
|
||||
else:
|
||||
rows = await AnalysisRepository.list_results(
|
||||
conn,
|
||||
args.run_id,
|
||||
args.result_type,
|
||||
limit=args.limit,
|
||||
lookahead=True,
|
||||
offset=args.offset,
|
||||
)
|
||||
data, page, truncated = page_result(rows, args)
|
||||
elif name in {"list_sensor_placement_runs", "get_sensor_placement_run"}:
|
||||
rows = await reads.sensor_runs(
|
||||
conn,
|
||||
run_id=getattr(args, "run_id", None),
|
||||
limit=getattr(args, "limit", 1) + 1,
|
||||
offset=getattr(args, "offset", 0),
|
||||
)
|
||||
if name == "get_sensor_placement_run":
|
||||
if not rows:
|
||||
raise HTTPException(404, "Sensor placement run was not found")
|
||||
data = rows[0]
|
||||
else:
|
||||
data, page, truncated = page_result(rows, args)
|
||||
elif name == "query_element_history":
|
||||
from app.domain.schemas.timeseries_history import ElementHistoryQuery
|
||||
from app.services.timeseries_history import TimeseriesHistoryService
|
||||
|
||||
query = ElementHistoryQuery(
|
||||
start_time=args.start,
|
||||
end_time=args.end,
|
||||
mode=args.mode,
|
||||
run_id=args.run_id,
|
||||
elements=[e.model_dump() for e in args.elements],
|
||||
)
|
||||
async with scope.timescale() as ts, conn.transaction(), ts.transaction():
|
||||
result = await TimeseriesHistoryService.query(
|
||||
ts, conn, query, row_limit=args.limit, row_offset=args.offset
|
||||
)
|
||||
data = result.model_dump(mode="json")
|
||||
truncated = result.truncated
|
||||
page = PageInfo(limit=args.limit, offset=args.offset, has_more=truncated)
|
||||
truncated = truncated or args.offset > 0
|
||||
elif name in {"get_element_gis_properties", "get_source_gis_properties"}:
|
||||
async with conn.transaction():
|
||||
if name == "get_element_gis_properties":
|
||||
data = await GisAttributeRepository.list_element_groups(
|
||||
conn, args.element_type, args.element_id
|
||||
)
|
||||
else:
|
||||
data = await GisAttributeRepository.get_source_feature_group(
|
||||
conn, args.layer_key, args.source_fid
|
||||
)
|
||||
if data is None:
|
||||
raise HTTPException(404, "Source GIS feature was not found")
|
||||
elif name == "get_task_context":
|
||||
if args.section in {"summary", "skill"}:
|
||||
data = await reads.task_context_summary(conn, project_id=binding.project_id,
|
||||
task_id=binding.task_id, report_revision=args.report_revision)
|
||||
skill = read_work_order_skill(
|
||||
data.get("task", {}).get("workflow"), include_content=args.section == "skill"
|
||||
)
|
||||
if args.section == "skill":
|
||||
data = {"section": "skill", "skill": skill}
|
||||
else:
|
||||
data["handling_skill"] = skill
|
||||
data["display_time_zone"] = "Asia/Shanghai"
|
||||
data["approved_data_scopes"] = binding.task_scope_cap
|
||||
data["bound_run_id"] = str(binding.run_id) if binding.run_id else None
|
||||
data["available_sections"] = ["task", "report", "evidence", "actions", "progress", "access_requests", "plan", "run", "skill"]
|
||||
truncated = True
|
||||
gaps.append("Task context is a bounded summary; use section, data_path, limit and offset for detail. Previews are not complete reports.")
|
||||
else:
|
||||
projection = await reads.task_context_page(conn, project_id=binding.project_id, task_id=binding.task_id,
|
||||
section=args.section, report_revision=args.report_revision, data_path=args.data_path, limit=args.limit, offset=args.offset)
|
||||
total = projection["total"]
|
||||
if args.section in {"evidence", "actions", "access_requests"} and not args.data_path:
|
||||
total = projection["snapshot_metadata"].get("total", total)
|
||||
if projection["value_type"] in {"array", "object", "string"}:
|
||||
page = PageInfo(limit=args.limit, offset=args.offset, total=total, has_more=args.offset + args.limit < total)
|
||||
elif args.offset:
|
||||
raise HTTPException(422, "Scalar task detail cannot use an offset")
|
||||
truncated = bool(projection.get("nested") or args.offset or page and page.has_more)
|
||||
data = {"section": args.section, "selected_path": args.data_path, "value": projection["value"], "metadata": projection["snapshot_metadata"]}
|
||||
if projection.get("nested"):
|
||||
gaps.append("Nested values have data_path descriptors; select a descriptor path for its values")
|
||||
elif name == "get_task_evidence":
|
||||
projection = await reads.evidence_page(
|
||||
conn, project_id=scope.project.project_id, task_id=binding.task_id,
|
||||
evidence_id=args.evidence_id, data_path=args.data_path,
|
||||
limit=args.limit, offset=args.offset, mcp_sources=list(TOOLS),
|
||||
full_snapshot=full_snapshot,
|
||||
)
|
||||
document = projection if full_snapshot else projection["snapshot_metadata"]
|
||||
origin = document["origin"]
|
||||
source_ref = document["source_ref"]
|
||||
evidence_id = args.evidence_id
|
||||
truncated = document.get("truncated", False)
|
||||
if full_snapshot:
|
||||
data = document
|
||||
else:
|
||||
selected = projection["value"]
|
||||
if projection["value_type"] in {"array", "object", "string"}:
|
||||
total = projection["total"]
|
||||
page = PageInfo(
|
||||
limit=args.limit,
|
||||
offset=args.offset,
|
||||
total=total,
|
||||
has_more=args.offset + args.limit < total,
|
||||
)
|
||||
truncated = truncated or page.has_more or args.offset > 0
|
||||
elif args.offset:
|
||||
raise HTTPException(422, "Scalar evidence cannot use an offset")
|
||||
if projection.get("nested"):
|
||||
truncated = True
|
||||
gaps.append(
|
||||
"Nested collections have data_path descriptors; query a descriptor path for their values"
|
||||
)
|
||||
data = {
|
||||
"snapshot_metadata": document,
|
||||
"query_metadata": projection["query_metadata"],
|
||||
"selected_path": args.data_path,
|
||||
"value": selected,
|
||||
}
|
||||
elif name == "get_workbench_runtime":
|
||||
from app.services.workbench_runtime import read_workbench_runtime
|
||||
|
||||
data = (await read_workbench_runtime(scope.project, conn)).model_dump(
|
||||
mode="json"
|
||||
)
|
||||
elif name == "query_peripheral_data":
|
||||
from app.mcp.peripheral import read_peripheral_query
|
||||
|
||||
snapshot = await read_peripheral_query(scope, args)
|
||||
data = snapshot.data
|
||||
origin = snapshot.origin
|
||||
truncated = snapshot.truncated
|
||||
else:
|
||||
raise ValueError("query handler is not registered")
|
||||
window = (
|
||||
{"start": args.start, "end": args.end}
|
||||
if hasattr(args, "start") and args.start is not None
|
||||
else None
|
||||
)
|
||||
data, nonfinite = finite_data(jsonable_encoder(data))
|
||||
semantics = (
|
||||
"comparison"
|
||||
if name == "query_element_history" and args.mode != "observed"
|
||||
else tool.semantics
|
||||
)
|
||||
return QueryResult(
|
||||
source=name,
|
||||
source_ref=source_ref,
|
||||
evidence_id=evidence_id,
|
||||
queried_at=datetime.now(timezone.utc),
|
||||
semantics=semantics,
|
||||
data=data,
|
||||
data_gaps=gaps
|
||||
+ (
|
||||
[f"{nonfinite} non-finite numeric values were returned as null"]
|
||||
if nonfinite
|
||||
else []
|
||||
),
|
||||
page=page,
|
||||
window=window,
|
||||
units=units,
|
||||
truncated=truncated,
|
||||
origin=origin,
|
||||
)
|
||||
@@ -0,0 +1,70 @@
|
||||
"""MCP online recalculation and isolated scenario adapters."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
SCENARIO_TOOLS = frozenset({
|
||||
"run_simulation",
|
||||
"run_contaminant_simulation", "analyze_valve_closure",
|
||||
"analyze_flushing", "analyze_pressure_regulation",
|
||||
})
|
||||
|
||||
|
||||
def calculate_scenario(name: str, parameters: dict[str, Any], *, network: str, username: str) -> dict[str, Any]:
|
||||
from app.services import simulation_scenarios as scenarios
|
||||
|
||||
parameters = dict(parameters)
|
||||
nodes = parameters.pop("output_node_ids")
|
||||
links = parameters.pop("output_link_ids")
|
||||
start = parameters.pop("start_time")
|
||||
duration = parameters.pop("duration")
|
||||
common = dict(name=network, modify_pattern_start_time=start.isoformat(),
|
||||
modify_total_duration=duration, persist_results=False, scheme_name=None)
|
||||
if name == "run_simulation":
|
||||
from app.services.simulation import run_online_simulation
|
||||
|
||||
result = run_online_simulation(
|
||||
network, start, duration, output_node_ids=nodes, output_link_ids=links,
|
||||
)
|
||||
elif name == "run_contaminant_simulation":
|
||||
result = scenarios.contaminant_simulation(**common, username=username,
|
||||
source=parameters["source"], concentration=parameters["concentration"], source_pattern=parameters.get("pattern"))
|
||||
elif name == "analyze_valve_closure":
|
||||
result = scenarios.valve_close_analysis(**common,
|
||||
modify_valve_opening={valve: 0.0 for valve in parameters["valves"]})
|
||||
elif name == "analyze_flushing":
|
||||
# The service consumes m³/h (the older REST label incorrectly says L/s).
|
||||
# Reject unsupported conversion rather than silently ignore the flow.
|
||||
from app.native.wndb.model.options import get_option
|
||||
if get_option(network).get("UNITS") not in {"LPS", "CMH"}:
|
||||
raise ValueError("Flushing currently supports LPS or CMH models only")
|
||||
result = 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"]})
|
||||
elif name == "analyze_pressure_regulation":
|
||||
from app.services.simulation import query_corresponding_element_id_and_query_id
|
||||
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("Pump controls must reference registered model pumps")
|
||||
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-speed pump controls must be 0 or 1")
|
||||
result = scenarios.pressure_regulation(**common,
|
||||
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)
|
||||
else:
|
||||
raise ValueError("Scenario is not registered")
|
||||
if not isinstance(result, dict):
|
||||
raise ValueError("Scenario did not return calculation results")
|
||||
selected_nodes = [row for row in result.get("node_results", []) if str(row.get("node")) in nodes]
|
||||
selected_links = [row for row in result.get("link_results", []) if str(row.get("link")) in links]
|
||||
return {
|
||||
"input_source": "online_model_simulation" if name == "run_simulation" else "isolated_model_simulation",
|
||||
"persisted_baseline": name == "run_simulation", "persisted_scheme": False,
|
||||
"start_time": start.isoformat(), "duration_seconds": duration,
|
||||
"units": result.get("units", {}), "times": result.get("times", {}),
|
||||
"node_results": selected_nodes, "link_results": selected_links,
|
||||
"coverage": {"requested_nodes": nodes, "requested_links": links,
|
||||
"missing_nodes": sorted(set(nodes) - {str(row["node"]) for row in selected_nodes}),
|
||||
"missing_links": sorted(set(links) - {str(row["link"]) for row in selected_links})},
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
"""Resolve authoritative metadata and project routing for each MCP request."""
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, AsyncIterator
|
||||
from psycopg import AsyncConnection
|
||||
from app.auth.project_dependencies import ProjectContext
|
||||
from app.infra.db.project_routing import ActiveProjectRouting
|
||||
from fastapi import HTTPException
|
||||
from mcp.server.auth.provider import AccessToken
|
||||
|
||||
from app.auth.metadata_dependencies import get_current_metadata_user
|
||||
from app.auth.project_dependencies import (
|
||||
resolve_project_context,
|
||||
resolve_project_business_routing,
|
||||
resolve_project_simulation_routing,
|
||||
get_project_pg_connection,
|
||||
get_project_timescale_connection,
|
||||
)
|
||||
from app.infra.db.metadb.database import get_metadata_session
|
||||
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
|
||||
from app.infra.db.project_routing import activate_project_routing
|
||||
|
||||
|
||||
class QueryMetadataRepository(MetadataRepository):
|
||||
async def get_project_db_routing(self, project_id, role):
|
||||
routing = await super().get_project_db_routing(project_id, role)
|
||||
# Existing connection generators obtain routing before awaiting a pool.
|
||||
# End that implicit SELECT transaction here, before the network wait.
|
||||
await self.session.commit()
|
||||
return routing
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueryScope:
|
||||
project: ProjectContext
|
||||
user: Any
|
||||
connection: AsyncConnection
|
||||
routing: ActiveProjectRouting
|
||||
metadata: QueryMetadataRepository
|
||||
|
||||
@asynccontextmanager
|
||||
async def simulation_routing(self) -> AsyncIterator[ActiveProjectRouting]:
|
||||
"""Resolve both databases before a native calculation can modify its model."""
|
||||
routing = await resolve_project_simulation_routing(self.project, self.metadata)
|
||||
await self.metadata.session.commit()
|
||||
with activate_project_routing(routing):
|
||||
yield routing
|
||||
|
||||
@asynccontextmanager
|
||||
async def timescale(self):
|
||||
# Metadata repository here is used only to resolve a known project DB.
|
||||
async with asynccontextmanager(get_project_timescale_connection)(
|
||||
self.project, self.metadata
|
||||
) as conn:
|
||||
await self.metadata.session.commit()
|
||||
yield conn
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def project_scope(principal: AccessToken, project_id: str):
|
||||
if not principal.claims:
|
||||
raise HTTPException(401, "MCP identity is missing")
|
||||
async with asynccontextmanager(get_metadata_session)() as session:
|
||||
metadata = QueryMetadataRepository(session)
|
||||
user = await get_current_metadata_user(principal.claims, metadata)
|
||||
project = await resolve_project_context(project_id, user, metadata)
|
||||
routing = await resolve_project_business_routing(project, metadata)
|
||||
# End the metadata read before any external query/IO waits.
|
||||
await session.commit()
|
||||
async with asynccontextmanager(get_project_pg_connection)(
|
||||
project, metadata
|
||||
) as conn:
|
||||
await session.commit()
|
||||
with activate_project_routing(routing):
|
||||
yield QueryScope(project, user, conn, routing, metadata)
|
||||
@@ -0,0 +1,40 @@
|
||||
"""Server-owned business instructions, selected from the persisted task workflow."""
|
||||
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
SKILL_ROOT = Path(__file__).resolve().parents[2] / "resources" / "work-order-skills"
|
||||
WORKFLOW_SKILLS = {
|
||||
"burst_report": "burst-report",
|
||||
"water_quality_complaint": "water-quality-complaint",
|
||||
"general": "appointment-follow-up",
|
||||
}
|
||||
|
||||
|
||||
def read_work_order_skill(workflow: str | None, *, include_content: bool = False) -> dict[str, Any]:
|
||||
"""Read only packaged, allowlisted instructions; never interpret a task path."""
|
||||
name = WORKFLOW_SKILLS.get(workflow)
|
||||
if name is None:
|
||||
return {"available": False, "workflow": workflow, "reason": "No business skill is configured for this workflow"}
|
||||
content = (SKILL_ROOT / name / "SKILL.md").read_text(encoding="utf-8")
|
||||
_, frontmatter, body = content.split("---", 2)
|
||||
header = yaml.safe_load(frontmatter)
|
||||
if header["name"] != name:
|
||||
raise ValueError("Packaged work-order skill name does not match its directory")
|
||||
result = {
|
||||
"available": True,
|
||||
"kind": "business_instructions",
|
||||
"workflow": workflow,
|
||||
"name": header["name"],
|
||||
"description": header["description"],
|
||||
"version": str(header["metadata"]["version"]),
|
||||
"content_sha256": sha256(content.encode("utf-8")).hexdigest(),
|
||||
"read": {"section": "skill"},
|
||||
}
|
||||
if include_content:
|
||||
result["instructions"] = body.strip()
|
||||
return result
|
||||
Reference in New Issue
Block a user