518 lines
21 KiB
Python
518 lines
21 KiB
Python
"""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
|