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())