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.
901 lines
31 KiB
Python
901 lines
31 KiB
Python
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())
|