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,900 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from fastapi import HTTPException
|
||||
from jose import jwt
|
||||
from mcp import types
|
||||
from mcp.server.auth.provider import AccessToken
|
||||
from starlette.applications import Starlette
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from app.mcp.application import McpQueryApplication
|
||||
from app.mcp.auth import McpTokenVerifier
|
||||
from app.mcp.config import McpSettings
|
||||
from app.mcp.catalog import TOOLS
|
||||
from app.mcp.models import BindInput, BoundContext, QueryResult, ScadaInput
|
||||
from app.mcp import application, contexts, query_service
|
||||
|
||||
PROJECT = uuid4()
|
||||
SUBJECT = str(uuid4())
|
||||
PERMISSIONS = {
|
||||
"webgis.view",
|
||||
"simulation.view",
|
||||
"scada.view",
|
||||
"operational_task.view",
|
||||
"operational_task.execute",
|
||||
"optimization.run",
|
||||
}
|
||||
|
||||
|
||||
def test_calculation_budget_outlives_query_default_and_still_times_out(runtime, monkeypatch):
|
||||
app, client, _, _ = runtime
|
||||
app.settings.query_timeout_seconds = 1
|
||||
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
||||
|
||||
async def execute(scope, bound, tool, args):
|
||||
await asyncio.sleep(1.05)
|
||||
return QueryResult(
|
||||
source=tool.name, source_ref="test:calculation", semantics="analysis",
|
||||
queried_at=datetime.now(timezone.utc), data={"completed": True},
|
||||
)
|
||||
|
||||
app.executor = execute
|
||||
handle = binding().context_handle
|
||||
arguments = {
|
||||
"context_handle": handle, "start_time": "2026-10-08T12:00:00+08:00",
|
||||
"duration": 3600, "output_node_ids": ["J1"], "timeout_seconds": 2,
|
||||
}
|
||||
response = request(client, "tools/call", {"name": "run_simulation", "arguments": arguments}).json()
|
||||
assert response["result"].get("isError") is not True
|
||||
assert response["result"]["structuredContent"]["data"] == {"completed": True}
|
||||
arguments["timeout_seconds"] = 1
|
||||
response = request(client, "tools/call", {"name": "run_simulation", "arguments": arguments}).json()
|
||||
assert response["result"]["isError"] is True
|
||||
assert response["result"]["structuredContent"]["error"]["code"] == "QUERY_TIMEOUT"
|
||||
|
||||
|
||||
def test_calculation_limit_is_advertised_and_enforced_before_execution(runtime):
|
||||
app, client, calls, _ = runtime
|
||||
app.settings.max_calculation_timeout_seconds = 120
|
||||
tools = request(client).json()["result"]["tools"]
|
||||
schema = next(tool["inputSchema"] for tool in tools if tool["name"] == "run_simulation")
|
||||
assert schema["properties"]["timeout_seconds"]["maximum"] == 120
|
||||
assert schema["properties"]["timeout_seconds"]["default"] == 120
|
||||
arguments = {
|
||||
"context_handle": binding().context_handle, "start_time": "2026-10-08T12:00:00+08:00",
|
||||
"duration": 3600, "output_node_ids": ["J1"], "timeout_seconds": 121,
|
||||
}
|
||||
response = request(client, "tools/call", {"name": "run_simulation", "arguments": arguments}).json()
|
||||
assert "error" in response or response["result"]["isError"] is True
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_burst_location_default_inputs_pass_protocol_validation(runtime, monkeypatch):
|
||||
app, client, calls, _ = runtime
|
||||
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
||||
schema = next(tool["inputSchema"] for tool in request(client).json()["result"]["tools"]
|
||||
if tool["name"] == "run_burst_location")
|
||||
assert "pressure_scada_ids" not in schema["required"]
|
||||
response = request(client, "tools/call", {"name": "run_burst_location", "arguments": {
|
||||
"context_handle": binding().context_handle, "burst_leakage": 0.0274,
|
||||
"scada_burst_start": "2026-10-08T12:00:00+08:00",
|
||||
"scada_burst_end": "2026-10-08T13:00:00+08:00",
|
||||
}}).json()
|
||||
assert response["result"].get("isError") is not True
|
||||
assert calls == [("session-1", "run_burst_location")]
|
||||
|
||||
|
||||
def settings():
|
||||
return McpSettings(
|
||||
enabled=True,
|
||||
issuer_uri="http://localhost/realms/water",
|
||||
resource_uri="http://localhost/mcp",
|
||||
allowed_hosts=["testserver", "localhost:*"],
|
||||
_env_file=None,
|
||||
)
|
||||
|
||||
|
||||
def principal(subject=SUBJECT, scopes=None, client="tjwater-agent-mcp"):
|
||||
return AccessToken(
|
||||
token="fake",
|
||||
subject=subject,
|
||||
client_id=client,
|
||||
scopes=scopes
|
||||
if scopes is not None
|
||||
else ["mcp.catalog.read", "mcp.query", "mcp.context.bind"],
|
||||
resource="http://localhost/mcp",
|
||||
claims={"sub": subject},
|
||||
)
|
||||
|
||||
|
||||
class Verifier:
|
||||
def __init__(self):
|
||||
self.identities = {
|
||||
"user": principal(),
|
||||
"other": principal(str(uuid4())),
|
||||
"catalog": principal(scopes=["mcp.catalog.read"]),
|
||||
}
|
||||
|
||||
async def verify_token(self, token):
|
||||
return self.identities.get(token)
|
||||
|
||||
|
||||
class Connection:
|
||||
@asynccontextmanager
|
||||
async def transaction(self):
|
||||
yield
|
||||
|
||||
|
||||
class Repository:
|
||||
bindings = {}
|
||||
receipts = {}
|
||||
|
||||
def __init__(self, conn, project):
|
||||
self.project = project
|
||||
|
||||
async def create(self, binding):
|
||||
self.bindings[binding.context_handle] = binding.model_copy(deep=True)
|
||||
|
||||
async def get(self, handle):
|
||||
return self.bindings.get(handle)
|
||||
|
||||
async def revoke(self, handle, subject, client):
|
||||
binding = self.bindings.get(handle)
|
||||
if (
|
||||
binding
|
||||
and binding.subject == subject
|
||||
and binding.client_id == client
|
||||
and binding.project_id == self.project
|
||||
):
|
||||
binding.revoked_at = datetime.now(timezone.utc)
|
||||
|
||||
async def receipt(self, key, call):
|
||||
return self.receipts.get((self.project, key, call))
|
||||
|
||||
async def save_receipt(self, key, call, parameters, result):
|
||||
self.receipts[(self.project, key, call)] = {
|
||||
"parameters": parameters,
|
||||
"result": result,
|
||||
}
|
||||
|
||||
async def lock_receipt(self, key, call):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runtime(monkeypatch):
|
||||
Repository.bindings = {}
|
||||
Repository.receipts = {}
|
||||
monkeypatch.setattr(application, "McpQueryRepository", Repository)
|
||||
monkeypatch.setattr(contexts, "McpQueryRepository", Repository)
|
||||
monkeypatch.setattr(contexts, "permissions_for_context", lambda ctx: PERMISSIONS)
|
||||
monkeypatch.setattr(
|
||||
query_service, "permissions_for_context", lambda ctx: PERMISSIONS
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def scope_factory(identity, project):
|
||||
try:
|
||||
project_id = __import__("uuid").UUID(project)
|
||||
except ValueError:
|
||||
raise HTTPException(400, "Invalid project")
|
||||
yield SimpleNamespace(
|
||||
project=SimpleNamespace(project_id=project_id), connection=Connection()
|
||||
)
|
||||
|
||||
calls = []
|
||||
|
||||
async def execute(scope, binding, tool, args):
|
||||
calls.append((binding.native_session_id, tool.name))
|
||||
return QueryResult(
|
||||
source=tool.name,
|
||||
source_ref="source:1",
|
||||
queried_at=datetime.now(timezone.utc),
|
||||
semantics=tool.semantics,
|
||||
data=[{"id": "P-1"}],
|
||||
)
|
||||
|
||||
verifier = Verifier()
|
||||
runtime = McpQueryApplication(
|
||||
settings(), verifier=verifier, scope_factory=scope_factory, executor=execute
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app):
|
||||
async with runtime.lifespan():
|
||||
yield
|
||||
|
||||
host = Starlette(lifespan=lifespan)
|
||||
|
||||
# Lifespan belongs to host; HTTP is delegated to SDK without mounting path prefixes.
|
||||
class App:
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] == "lifespan":
|
||||
await host(scope, receive, send)
|
||||
else:
|
||||
await runtime.asgi(scope, receive, send)
|
||||
|
||||
with TestClient(App()) as client:
|
||||
yield runtime, client, calls, verifier
|
||||
|
||||
|
||||
def request(
|
||||
client,
|
||||
method="tools/list",
|
||||
params=None,
|
||||
token="user",
|
||||
version="2026-07-28",
|
||||
extra=None,
|
||||
):
|
||||
meta = {
|
||||
types.PROTOCOL_VERSION_META_KEY: version,
|
||||
types.CLIENT_CAPABILITIES_META_KEY: {},
|
||||
types.CLIENT_INFO_META_KEY: {"name": "test", "version": "1"},
|
||||
"io.tjwater/callId": "call-1",
|
||||
}
|
||||
body = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": method,
|
||||
"params": {"_meta": meta, **(params or {})},
|
||||
}
|
||||
headers = {
|
||||
"Authorization": "Bearer " + token,
|
||||
"MCP-Protocol-Version": version,
|
||||
"MCP-Method": method,
|
||||
"Accept": "application/json, text/event-stream",
|
||||
"X-Project-Id": str(PROJECT),
|
||||
}
|
||||
if method == "tools/call":
|
||||
headers["MCP-Name"] = params["name"]
|
||||
headers.update(extra or {})
|
||||
return client.post("/mcp", json=body, headers=headers)
|
||||
|
||||
|
||||
def binding(task=False):
|
||||
value = BoundContext(
|
||||
context_handle="a" * 43,
|
||||
project_id=PROJECT,
|
||||
subject=SUBJECT,
|
||||
client_id="tjwater-agent-mcp",
|
||||
native_session_id="session-1",
|
||||
purpose="task_execution" if task else "ordinary_chat",
|
||||
tool_cap=sorted(TOOLS),
|
||||
permission_cap=sorted(PERMISSIONS),
|
||||
task_id=uuid4() if task else None,
|
||||
run_id=uuid4() if task else None,
|
||||
expires_at=datetime.now(timezone.utc) + timedelta(minutes=5),
|
||||
)
|
||||
Repository.bindings[value.context_handle] = value
|
||||
return value
|
||||
|
||||
|
||||
def test_modern_discovery_catalog_and_metadata(runtime):
|
||||
_, client, _, _ = runtime
|
||||
result = request(client, "server/discover").json()
|
||||
assert result["result"]["resultType"] == "complete"
|
||||
assert "2026-07-28" in result["result"]["supportedVersions"]
|
||||
listed = request(client).json()["result"]
|
||||
assert len(listed["tools"]) == 35
|
||||
for tool in listed["tools"]:
|
||||
timeout = tool["inputSchema"]["properties"].get("timeout_seconds")
|
||||
assert (timeout is not None) == TOOLS[tool["name"]].computation
|
||||
if timeout:
|
||||
assert timeout["default"] == 120
|
||||
assert timeout["maximum"] == 3600
|
||||
assert listed["ttlMs"] == 300000 and listed["cacheScope"] == "public"
|
||||
assert all(
|
||||
tool["inputSchema"]["additionalProperties"] is False for tool in listed["tools"]
|
||||
)
|
||||
assert (
|
||||
client.get("/.well-known/oauth-protected-resource/mcp").json()["resource"]
|
||||
== "http://localhost/mcp"
|
||||
)
|
||||
assert "mcp-session-id" not in request(client).headers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("token,status", [("bad", 401), ("catalog", 403)])
|
||||
def test_oauth_challenges(runtime, token, status):
|
||||
_, client, _, _ = runtime
|
||||
result = request(
|
||||
client, "tools/call", {"name": "list_elements", "arguments": {}}, token=token
|
||||
)
|
||||
assert result.status_code == status
|
||||
assert "resource_metadata=" in result.headers["www-authenticate"]
|
||||
assert "mcp.query" in result.headers["www-authenticate"]
|
||||
|
||||
|
||||
def test_rejects_legacy_and_mirror_mismatch(runtime):
|
||||
_, client, _, _ = runtime
|
||||
assert request(client, version="2025-03-26").json()["error"]["code"] == -32022
|
||||
assert client.get("/mcp").status_code == 405
|
||||
assert client.delete("/mcp").status_code == 405
|
||||
mismatch = request(client, extra={"MCP-Method": "resources/list"})
|
||||
assert mismatch.status_code == 400 and mismatch.json()["error"]["code"] == -32020
|
||||
unknown = request(client, "tools/call", {"name": "missing", "arguments": {}})
|
||||
assert unknown.json()["error"]["code"] == -32602
|
||||
|
||||
|
||||
@pytest.mark.parametrize("header,body", [
|
||||
("2025-11-25", "2026-07-28"),
|
||||
("1900-01-01", "2026-07-28"),
|
||||
("2026-07-28", "2025-11-25"),
|
||||
])
|
||||
def test_version_header_mismatch_precedes_version_support(runtime, header, body):
|
||||
_, client, calls, _ = runtime
|
||||
response = request(client, version=body, extra={"MCP-Protocol-Version": header})
|
||||
assert response.status_code == 400
|
||||
assert response.json()["error"]["code"] == -32020
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_query_identity_project_expiry_and_permission_boundaries(runtime):
|
||||
_, client, calls, _ = runtime
|
||||
value = binding()
|
||||
params = {
|
||||
"name": "list_elements",
|
||||
"arguments": {"context_handle": value.context_handle, "element_type": "pipe"},
|
||||
}
|
||||
assert request(client, "tools/call", params).json()["result"]["structuredContent"][
|
||||
"data"
|
||||
] == [{"id": "P-1"}]
|
||||
assert len(calls) == 1
|
||||
for token, extra in [("other", None), ("user", {"X-Project-Id": str(uuid4())})]:
|
||||
assert request(client, "tools/call", params, token=token, extra=extra).json()[
|
||||
"result"
|
||||
]["isError"]
|
||||
value.permission_cap = []
|
||||
assert (
|
||||
request(client, "tools/call", params).json()["result"]["structuredContent"][
|
||||
"error"
|
||||
]["code"]
|
||||
== "QUERY_CONTEXT_DENIED"
|
||||
)
|
||||
value.permission_cap = sorted(PERMISSIONS)
|
||||
value.expires_at = datetime.now(timezone.utc) - timedelta(seconds=1)
|
||||
assert request(client, "tools/call", params).json()["result"]["isError"]
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_host_only_binding_and_no_identity_arguments():
|
||||
with pytest.raises(Exception):
|
||||
BindInput(native_session_id="s", purpose="ordinary_chat", task_id=uuid4())
|
||||
with pytest.raises(Exception):
|
||||
BindInput(native_session_id="s", purpose="task_reader", task_id=uuid4())
|
||||
with pytest.raises(Exception):
|
||||
ScadaInput(
|
||||
context_handle="x" * 43,
|
||||
device_ids=["d"],
|
||||
start="2026-10-01T00:00:00",
|
||||
end="2026-10-02T00:00:00",
|
||||
)
|
||||
with pytest.raises(Exception):
|
||||
TOOLS["list_elements"].model.model_validate(
|
||||
{
|
||||
"context_handle": "x" * 43,
|
||||
"element_type": "pipe",
|
||||
"project_id": str(PROJECT),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_small_page_defaults_and_invalid_deep_offset(runtime):
|
||||
_, client, calls, _ = runtime
|
||||
value = binding()
|
||||
for name in ("list_elements", "list_analysis_runs", "query_scada_readings", "query_model_baseline_field"):
|
||||
assert TOOLS[name].model.model_fields["limit"].default == 20
|
||||
result = request(client, "tools/call", {
|
||||
"name": "list_elements", "arguments": {
|
||||
"context_handle": value.context_handle, "element_type": "pipe", "offset": 10001,
|
||||
},
|
||||
}).json()["result"]
|
||||
assert result["isError"] and calls == []
|
||||
|
||||
|
||||
def test_receipt_retry_checks_authorization_before_returning_cache(
|
||||
runtime, monkeypatch
|
||||
):
|
||||
_, client, calls, _ = runtime
|
||||
value = binding(task=True)
|
||||
|
||||
async def verify(scope, binding):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(contexts, "verify_task", verify)
|
||||
# Existing persisted receipt, not a second execution/evidence append.
|
||||
import json
|
||||
|
||||
key = json.dumps(
|
||||
[
|
||||
value.subject,
|
||||
value.client_id,
|
||||
value.native_session_id,
|
||||
str(value.task_id),
|
||||
str(value.run_id),
|
||||
str(value.execution_request_id),
|
||||
],
|
||||
separators=(",", ":"),
|
||||
)
|
||||
params = {
|
||||
"tool": "list_elements",
|
||||
"arguments": {"element_type": "pipe", "limit": 20, "offset": 0},
|
||||
}
|
||||
result = QueryResult(
|
||||
source="list_elements",
|
||||
source_ref="mcp:saved",
|
||||
queried_at=datetime.now(timezone.utc),
|
||||
semantics="asset",
|
||||
data=[],
|
||||
evidence_id="evidence:1",
|
||||
).model_dump(mode="json")
|
||||
Repository.receipts[(PROJECT, key, "call-1")] = {
|
||||
"parameters": params,
|
||||
"result": result,
|
||||
}
|
||||
query = {
|
||||
"name": "list_elements",
|
||||
"arguments": {"context_handle": value.context_handle, "element_type": "pipe"},
|
||||
}
|
||||
assert (
|
||||
request(client, "tools/call", query).json()["result"]["structuredContent"][
|
||||
"evidence_id"
|
||||
]
|
||||
== "evidence:1"
|
||||
)
|
||||
value.permission_cap = []
|
||||
assert request(client, "tools/call", query).json()["result"]["isError"]
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_evidence_is_saved_before_return_and_call_conflicts_fail(runtime, monkeypatch):
|
||||
_, client, calls, _ = runtime
|
||||
value = binding(task=True)
|
||||
saved = []
|
||||
|
||||
async def verify(scope, binding):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(contexts, "verify_task", verify)
|
||||
monkeypatch.setattr(application, "verify_task", verify)
|
||||
|
||||
class TaskService:
|
||||
repo = SimpleNamespace(lock=lambda key: asyncio.sleep(0))
|
||||
|
||||
async def record_data_snapshot(self, task_id, snapshot):
|
||||
saved.append(snapshot)
|
||||
return SimpleNamespace(evidence_id=snapshot.source_ref), snapshot
|
||||
|
||||
monkeypatch.setattr(application, "task_service", lambda scope: TaskService())
|
||||
query = {
|
||||
"name": "list_elements",
|
||||
"arguments": {"context_handle": value.context_handle, "element_type": "pipe"},
|
||||
}
|
||||
first = request(client, "tools/call", query).json()["result"]["structuredContent"]
|
||||
second = request(client, "tools/call", query).json()["result"]["structuredContent"]
|
||||
assert (
|
||||
first == second
|
||||
and first["evidence_id"] == saved[0].source_ref
|
||||
and len(saved) == len(calls) == 1
|
||||
)
|
||||
query["arguments"]["element_type"] = "junction"
|
||||
assert request(client, "tools/call", query).json()["result"]["isError"]
|
||||
assert len(saved) == 1
|
||||
|
||||
|
||||
def test_business_skill_is_delivered_over_mcp_without_recording_event_evidence(runtime, monkeypatch):
|
||||
app, client, _, _ = runtime
|
||||
value = binding(task=True)
|
||||
monkeypatch.setattr(contexts, "verify_task", AsyncMock())
|
||||
reader = AsyncMock(return_value={"task": {"workflow": "water_quality_complaint"}})
|
||||
monkeypatch.setattr(query_service.reads, "task_context_summary", reader)
|
||||
event_service = AsyncMock()
|
||||
monkeypatch.setattr(application, "task_service", event_service)
|
||||
app.executor = query_service.execute_query
|
||||
response = request(client, "tools/call", {
|
||||
"name": "get_task_context",
|
||||
"arguments": {"context_handle": value.context_handle, "section": "skill"},
|
||||
}).json()["result"]
|
||||
assert response.get("isError") is not True
|
||||
result = response["structuredContent"]
|
||||
assert result["data"]["skill"]["name"] == "water-quality-complaint"
|
||||
assert result["data"]["skill"]["instructions"] and not result["truncated"]
|
||||
assert result["evidence_id"] is None
|
||||
event_service.assert_not_called()
|
||||
assert not Repository.receipts
|
||||
|
||||
|
||||
def test_old_task_run_and_reader_parent_are_rejected(monkeypatch):
|
||||
monkeypatch.setattr(contexts, "permissions_for_context", lambda ctx: PERMISSIONS)
|
||||
value = binding(task=True)
|
||||
value.parent_session_id = "parent"
|
||||
run = {
|
||||
"status": "running",
|
||||
"agent_session_id": "wrong",
|
||||
"run_id": str(value.run_id),
|
||||
"execution_request_id": str(uuid4()),
|
||||
"lease_expires_at": datetime.now(timezone.utc) + timedelta(minutes=1),
|
||||
}
|
||||
|
||||
class Tasks:
|
||||
async def get_task(self, id):
|
||||
return {"task_id": id}
|
||||
|
||||
async def get_active_run(self, id):
|
||||
return run
|
||||
|
||||
monkeypatch.setattr(contexts, "task_repository", lambda scope: Tasks())
|
||||
scope = SimpleNamespace(project=object(), connection=Connection())
|
||||
with pytest.raises(HTTPException) as failure:
|
||||
asyncio.run(contexts.verify_task(scope, value))
|
||||
assert failure.value.status_code == 409
|
||||
run["agent_session_id"] = "parent"
|
||||
asyncio.run(contexts.verify_task(scope, value))
|
||||
run["run_id"] = str(uuid4())
|
||||
with pytest.raises(HTTPException):
|
||||
asyncio.run(contexts.verify_task(scope, value))
|
||||
|
||||
|
||||
def test_jwt_signature_issuer_audience_lifetime_and_key_cache():
|
||||
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
pem = key.private_bytes(
|
||||
serialization.Encoding.PEM,
|
||||
serialization.PrivateFormat.PKCS8,
|
||||
serialization.NoEncryption(),
|
||||
)
|
||||
public = key.public_key().public_numbers()
|
||||
|
||||
def encoded(number):
|
||||
return (
|
||||
base64.urlsafe_b64encode(
|
||||
number.to_bytes((number.bit_length() + 7) // 8, "big")
|
||||
)
|
||||
.rstrip(b"=")
|
||||
.decode()
|
||||
)
|
||||
|
||||
jwks = {
|
||||
"keys": [
|
||||
{
|
||||
"kid": "one",
|
||||
"kty": "RSA",
|
||||
"alg": "RS256",
|
||||
"n": encoded(public.n),
|
||||
"e": encoded(public.e),
|
||||
}
|
||||
]
|
||||
}
|
||||
calls = []
|
||||
|
||||
def respond(request):
|
||||
calls.append(request)
|
||||
return httpx.Response(200, json=jwks)
|
||||
|
||||
async def check():
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http:
|
||||
verifier = McpTokenVerifier(settings(), http)
|
||||
claims = {
|
||||
"sub": SUBJECT,
|
||||
"iss": settings().issuer_uri,
|
||||
"aud": "tjwater-mcp",
|
||||
"azp": "tjwater-agent-mcp",
|
||||
"iat": time.time(),
|
||||
"exp": time.time() + 300,
|
||||
"scope": "mcp.query mcp.catalog.read",
|
||||
}
|
||||
|
||||
def sign(overrides=None):
|
||||
return jwt.encode(
|
||||
{**claims, **(overrides or {})},
|
||||
pem,
|
||||
algorithm="RS256",
|
||||
headers={"kid": "one"},
|
||||
)
|
||||
|
||||
assert (await verifier.verify_token(sign())).subject == SUBJECT
|
||||
for overrides in [
|
||||
{"aud": "tjwater-api"},
|
||||
{"iss": "http://other"},
|
||||
{"exp": time.time() - 1},
|
||||
{"iat": time.time() + 120},
|
||||
{"sub": "not-a-user"},
|
||||
]:
|
||||
assert await verifier.verify_token(sign(overrides)) is None
|
||||
assert len(calls) == 1
|
||||
assert await verifier.verify_token("broken") is None
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_context_control_routes_are_host_only_and_support_revocation(
|
||||
runtime, monkeypatch
|
||||
):
|
||||
runtime, _, _, verifier = runtime
|
||||
from fastapi import FastAPI
|
||||
from app.api.v1.endpoints import mcp_queries
|
||||
|
||||
monkeypatch.setattr(mcp_queries, "McpQueryRepository", Repository)
|
||||
host = FastAPI()
|
||||
host.state.mcp_runtime = runtime
|
||||
host.include_router(runtime.router, prefix="/api/v1")
|
||||
command = {"native_session_id": "host-native", "purpose": "ordinary_chat"}
|
||||
headers = {"Authorization": "Bearer user", "X-Project-Id": str(PROJECT)}
|
||||
with TestClient(host) as client:
|
||||
result = client.post(
|
||||
"/api/v1/mcp-query-contexts", json=command, headers=headers
|
||||
)
|
||||
assert result.status_code == 200
|
||||
handle = result.json()["context_handle"]
|
||||
assert Repository.bindings[handle].permission_cap == sorted(PERMISSIONS)
|
||||
verifier.identities["outsider"] = principal(client="browser")
|
||||
denied = client.post(
|
||||
"/api/v1/mcp-query-contexts",
|
||||
json=command,
|
||||
headers={**headers, "Authorization": "Bearer outsider"},
|
||||
)
|
||||
assert denied.status_code == 403
|
||||
assert (
|
||||
client.post(
|
||||
"/api/v1/mcp-query-contexts",
|
||||
json={**command, "subject": SUBJECT},
|
||||
headers=headers,
|
||||
).status_code
|
||||
== 422
|
||||
)
|
||||
assert (
|
||||
client.delete(
|
||||
"/api/v1/mcp-query-contexts/" + handle,
|
||||
headers={**headers, "Authorization": "Bearer other"},
|
||||
).status_code
|
||||
== 204
|
||||
)
|
||||
assert Repository.bindings[handle].revoked_at is None
|
||||
assert (
|
||||
client.delete(
|
||||
"/api/v1/mcp-query-contexts/" + handle, headers=headers
|
||||
).status_code
|
||||
== 204
|
||||
)
|
||||
assert Repository.bindings[handle].revoked_at is not None
|
||||
host.state.mcp_runtime = None
|
||||
assert (
|
||||
client.post(
|
||||
"/api/v1/mcp-query-contexts", json=command, headers=headers
|
||||
).status_code
|
||||
== 503
|
||||
)
|
||||
|
||||
|
||||
def test_large_task_context_uses_bounded_summary_without_loading_full_context(runtime, monkeypatch):
|
||||
from unittest.mock import AsyncMock
|
||||
from app.infra.db.postgresql import read_queries
|
||||
|
||||
runtime, client, _, _ = runtime
|
||||
value = binding(task=True)
|
||||
monkeypatch.setattr(contexts, "verify_task", AsyncMock())
|
||||
full = SimpleNamespace(model_dump=lambda **kwargs: {"task": {"latest_progress": {"summary": "x" * 13000}}})
|
||||
service = SimpleNamespace(agent_context=AsyncMock(return_value=full))
|
||||
monkeypatch.setattr(query_service, "task_service", lambda scope: service)
|
||||
summary = AsyncMock(return_value={"task": {"task_id": str(value.task_id)}, "selected_report": {"revision": 3}})
|
||||
monkeypatch.setattr(read_queries, "task_context_summary", summary, raising=False)
|
||||
runtime.executor = query_service.execute_query
|
||||
response = request(client, "tools/call", {"name": "get_task_context", "arguments": {"context_handle": value.context_handle, "report_revision": 3}})
|
||||
result = response.json()["result"]
|
||||
assert response.status_code == 200 and not result.get("isError", False), result
|
||||
assert result["structuredContent"]["data"]["selected_report"]["revision"] == 3
|
||||
service.agent_context.assert_not_awaited()
|
||||
assert summary.await_args.kwargs["report_revision"] == 3
|
||||
|
||||
|
||||
def test_task_context_sections_support_projection_and_bounded_pages():
|
||||
from app.mcp.models import TaskContextInput
|
||||
|
||||
args = TaskContextInput(context_handle="a" * 43, section="report", report_revision=3,
|
||||
data_path=["report", "current_conclusion"], limit=100, offset=100)
|
||||
assert args.section == "report" and args.offset == 100
|
||||
with pytest.raises(ValueError):
|
||||
TaskContextInput(context_handle="a" * 43, section="evidence", offset=10001)
|
||||
|
||||
|
||||
def test_concurrency_bound_is_not_a_successful_empty_query(runtime):
|
||||
runtime, client, _, _ = runtime
|
||||
runtime.active[(SUBJECT, str(PROJECT))] = 4
|
||||
query = {
|
||||
"name": "list_elements",
|
||||
"arguments": {"context_handle": "a" * 43, "element_type": "pipe"},
|
||||
}
|
||||
assert (
|
||||
request(client, "tools/call", query).json()["result"]["structuredContent"][
|
||||
"error"
|
||||
]["code"]
|
||||
== "RATE_LIMITED"
|
||||
)
|
||||
assert runtime.active[(SUBJECT, str(PROJECT))] == 4
|
||||
|
||||
|
||||
def test_time_series_sql_is_bounded_and_field_names_cannot_inject_sql():
|
||||
from app.infra.db.postgresql import read_queries
|
||||
|
||||
statements = []
|
||||
|
||||
class Cursor:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
pass
|
||||
|
||||
async def execute(self, sql, params):
|
||||
statements.append(
|
||||
(sql.as_string() if hasattr(sql, "as_string") else sql, params)
|
||||
)
|
||||
|
||||
async def fetchall(self):
|
||||
return []
|
||||
|
||||
class Db(Connection):
|
||||
def cursor(self, **options):
|
||||
return Cursor()
|
||||
|
||||
start = datetime.now(timezone.utc)
|
||||
|
||||
async def check():
|
||||
await read_queries.series_rows(
|
||||
Db(),
|
||||
source="scada",
|
||||
element_type="junction",
|
||||
ids=["D'1"],
|
||||
start=start,
|
||||
end=start + timedelta(hours=1),
|
||||
field="monitored_value",
|
||||
limit=101,
|
||||
offset=200,
|
||||
)
|
||||
assert statements[-1][1][-2:] == (101, 200)
|
||||
assert "LIMIT %s OFFSET %s" in statements[-1][0]
|
||||
assert "D'1" not in statements[-1][0]
|
||||
assert "statement_timeout" in statements[0][0]
|
||||
with pytest.raises(ValueError):
|
||||
await read_queries.series_rows(
|
||||
Db(),
|
||||
source="realtime",
|
||||
element_type="pipe",
|
||||
ids=["P1"],
|
||||
time=start,
|
||||
field="pressure; DROP TABLE network.nodes",
|
||||
)
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_analysis_preserves_run_units_and_reports_inference(monkeypatch):
|
||||
from app.mcp.models import AnalysisValuesInput
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
monkeypatch.setattr(
|
||||
query_service, "permissions_for_context", lambda ctx: PERMISSIONS
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
query_service.reads, "series_rows", AsyncMock(return_value=[{"value": 3.5}])
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
query_service.NetworkSettingsRepository,
|
||||
"get_result_units",
|
||||
AsyncMock(return_value={"pressure": "m"}),
|
||||
)
|
||||
run = {"parameters": {"result_units": {"pressure": "bar"}}}
|
||||
monkeypatch.setattr(
|
||||
query_service.AnalysisRepository,
|
||||
"get_run",
|
||||
AsyncMock(side_effect=lambda *args: run),
|
||||
)
|
||||
|
||||
class Scope:
|
||||
project = object()
|
||||
connection = Connection()
|
||||
|
||||
@asynccontextmanager
|
||||
async def timescale(self):
|
||||
yield self.connection
|
||||
|
||||
args = AnalysisValuesInput(
|
||||
context_handle="a" * 43,
|
||||
run_id=uuid4(),
|
||||
element_type="junction",
|
||||
field="pressure",
|
||||
time=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
async def check():
|
||||
result = await query_service.execute_query(
|
||||
Scope(), binding(), TOOLS["query_analysis_values"], args
|
||||
)
|
||||
assert result.units == {"pressure": "bar"} and not result.data_gaps
|
||||
run["parameters"] = {}
|
||||
result = await query_service.execute_query(
|
||||
Scope(), binding(), TOOLS["query_analysis_values"], args
|
||||
)
|
||||
assert result.units == {"pressure": "m"} and "inferred" in result.data_gaps[0]
|
||||
assert query_service.finite_data([float("nan"), 2.0, float("inf")]) == (
|
||||
[None, 2.0, None],
|
||||
2,
|
||||
)
|
||||
|
||||
asyncio.run(check())
|
||||
|
||||
|
||||
def test_saved_evidence_pages_are_immutable_and_keep_synthetic_origin(monkeypatch):
|
||||
from app.domain.schemas.operational_data import TaskDataSnapshot
|
||||
from app.mcp.models import EvidenceInput
|
||||
|
||||
monkeypatch.setattr(
|
||||
query_service, "permissions_for_context", lambda ctx: PERMISSIONS
|
||||
)
|
||||
snapshot = TaskDataSnapshot(
|
||||
source="query_scada_readings",
|
||||
source_ref="evidence:stored",
|
||||
queried_at=datetime.now(timezone.utc),
|
||||
origin="synthetic",
|
||||
truncated=False,
|
||||
data={
|
||||
"semantics": "observed",
|
||||
"data": {
|
||||
"items": [{"value": 1}, {"value": 2}],
|
||||
"series": [{"points": [{"value": 3}, {"value": 4}]}],
|
||||
},
|
||||
},
|
||||
)
|
||||
async def project(conn, *, full_snapshot=False, **options):
|
||||
document = snapshot.model_dump(mode="json")
|
||||
if full_snapshot:
|
||||
return document
|
||||
selected = document["data"]["data"]
|
||||
for part in options["data_path"]:
|
||||
selected = selected[int(part)] if isinstance(selected, list) else selected[part]
|
||||
return {
|
||||
"snapshot_metadata": {k: v for k, v in document.items() if k != "data"},
|
||||
"query_metadata": {"semantics": "observed"},
|
||||
"value_type": "array", "total": len(selected),
|
||||
"value": selected[options["offset"]:options["offset"] + options["limit"]],
|
||||
}
|
||||
|
||||
monkeypatch.setattr(query_service.reads, "evidence_page", project)
|
||||
scope = SimpleNamespace(project=SimpleNamespace(project_id=PROJECT), connection=Connection())
|
||||
value = binding(task=True)
|
||||
|
||||
async def query(path, **options):
|
||||
args = EvidenceInput(
|
||||
context_handle=value.context_handle,
|
||||
evidence_id=snapshot.source_ref,
|
||||
data_path=path,
|
||||
limit=1,
|
||||
offset=1,
|
||||
)
|
||||
return await query_service.execute_query(
|
||||
scope, value, TOOLS["get_task_evidence"], args, **options
|
||||
)
|
||||
|
||||
async def check():
|
||||
page = await query(["items"])
|
||||
assert page.data["value"] == [{"value": 2}] and page.page.total == 2
|
||||
assert (
|
||||
page.origin == "synthetic"
|
||||
and page.evidence_id == page.source_ref == "evidence:stored"
|
||||
and page.truncated
|
||||
)
|
||||
assert (await query(["series", "0", "points"])).data["value"] == [{"value": 4}]
|
||||
full = await query([], full_snapshot=True)
|
||||
assert full.data == snapshot.model_dump(mode="json")
|
||||
assert snapshot.data["data"]["items"] == [{"value": 1}, {"value": 2}]
|
||||
|
||||
asyncio.run(check())
|
||||
Reference in New Issue
Block a user