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,505 @@
|
||||
"""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.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
|
||||
|
||||
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:
|
||||
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 else self.settings.query_timeout_seconds
|
||||
)
|
||||
if tool.computation 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 = TaskDataSnapshot(
|
||||
source=name,
|
||||
source_ref=result.source_ref,
|
||||
queried_at=result.queried_at,
|
||||
origin=result.origin,
|
||||
data=result.model_dump(mode="json"),
|
||||
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 = []
|
||||
if (
|
||||
not include_full
|
||||
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 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
|
||||
Reference in New Issue
Block a user