"""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.errors import CalculationCapacityExceeded 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 from app.domain.domain_errors import DomainError 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 or tool.mutating: 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 or tool.mutating else self.settings.query_timeout_seconds ) if (tool.computation or tool.mutating) 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_data = result.model_dump(mode="json") from app.mcp.extended_service import ARTIFACT_TOOLS if name in ARTIFACT_TOOLS: snapshot_data["data"]["file"].pop("data", None) snapshot = TaskDataSnapshot( source=name, source_ref=result.source_ref, queried_at=result.queried_at, origin=result.origin, data=snapshot_data, 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 = [] from app.mcp.extended_service import ARTIFACT_TOOLS if ( not include_full and name not in ARTIFACT_TOOLS 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 DomainError as exc: return tool_error(exc.code, exc.message) except CalculationCapacityExceeded as exc: return tool_error(exc.code, str(exc.detail)) 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