"""MCPServer instance exposing the compliance knowledge base as an MCP tool. This module is a pure protocol adapter: search_regulations() below calls the existing AgentConversationService.ask() (the same application service backing the /api/v1/agent/ask REST endpoint) and reshapes its result into a plain dict. No new retrieval, ranking, or LLM orchestration logic lives here. """ from __future__ import annotations import logging import time from typing import Annotated from mcp.server import MCPServer from mcp.server.transport_security import TransportSecuritySettings from pydantic import Field from starlette.responses import PlainTextResponse from starlette.types import ASGIApp, Receive, Scope, Send from app.config.settings import settings from app.mcp.stats import get_mcp_stats_tracker from app.shared.bootstrap import get_agent_conversation_service, get_jwt_handler logger = logging.getLogger(__name__) # Single shared MCPServer instance — analogous to the single shared FastAPI # `app` instance in app/api/main.py. Tools registered via @mcp.tool() below. # Note: the installed mcp SDK (2.0.0) renamed the older "FastMCP" class to # "MCPServer" (mcp.server.mcpserver.MCPServer); the .tool()/.streamable_http_app() # API surface used here is unchanged across that rename. mcp = MCPServer("ai-regulations") @mcp.tool() def search_regulations( query: Annotated[str, Field(min_length=1, max_length=2000)], top_k: Annotated[int, Field(ge=1, le=20)] = 5, ) -> dict: """Search the compliance knowledge base and return a grounded answer. query: Natural-language search question, e.g. "国六排放标准最新要求". top_k: Maximum number of cited sources to return (1-20, default 5). """ # Bounds mirror AskRequest in app/api/models/agent.py so the MCP path cannot # be used to bypass the REST endpoint's limits. They matter more here than # there: KnowledgeRetrievalService amplifies top_k (candidate_k = top_k * 4) # when reranking, and an LLM client can easily hallucinate a huge value. # Declaring them via Annotated puts them in the advertised JSON schema too, # so well-behaved clients never send an out-of-range value in the first place. # # No session_id is passed: this keeps each call stateless (no # ConversationStore reads/writes), matching "search" semantics rather # than multi-turn chat semantics. started = time.perf_counter() try: _, result = get_agent_conversation_service().ask(query=query, top_k=top_k) except Exception: # Record the failure, then re-raise unchanged so the MCP SDK still # converts it into a protocol-level error for the client. Swallowing # it here would report success to the caller. get_mcp_stats_tracker().record( tool="search_regulations", duration_ms=(time.perf_counter() - started) * 1000, success=False, ) raise get_mcp_stats_tracker().record( tool="search_regulations", duration_ms=(time.perf_counter() - started) * 1000, success=True, ) return { "answer": result.answer, "sources": [source.__dict__ for source in result.sources], } class MCPAuthMiddleware: """Reject unauthenticated requests before they reach the MCP protocol handler. Mirrors the existing get_current_user dependency's behavior (auth.py) but implemented as raw ASGI middleware, since the mounted MCP app is a plain ASGI app, not a FastAPI/APIRouter instance that supports Depends(). """ def __init__(self, app: ASGIApp) -> None: """Store the wrapped ASGI app to delegate to once auth passes.""" self.app = app async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: """Validate the bearer token for HTTP requests; pass non-HTTP scopes through.""" # Only HTTP requests carry an Authorization header to check; lifespan # and other scope types must always pass through untouched. if scope["type"] != "http" or not settings.auth_enabled: await self.app(scope, receive, send) return headers = dict(scope["headers"]) # ASGI header values are raw bytes specified as latin-1, not UTF-8; # decoding strictly as UTF-8 would raise on a malformed byte and turn a # bad request into an unhandled 500. auth_header = headers.get(b"authorization", b"").decode("latin-1") token = auth_header.removeprefix("Bearer ").strip() try: get_jwt_handler().decode_token(token) except ValueError as exc: # Reject before the MCP session/protocol layer ever sees the request. # WWW-Authenticate matches the get_current_user dependency (auth.py) # and is required by RFC 7235 so clients can tell "needs credentials" # apart from a generic failure. response = PlainTextResponse( str(exc), status_code=401, headers={"WWW-Authenticate": "Bearer"} ) await response(scope, receive, send) return await self.app(scope, receive, send) def _parse_allowed_hosts() -> list[str]: """Split the configured MCP host allow-list into individual entries.""" # Shared by the transport-security builder and the status endpoint so the # panel can never display an allow-list different from the enforced one. return [h.strip() for h in settings.mcp_allowed_hosts.split(",") if h.strip()] def _build_transport_security() -> TransportSecuritySettings: """Translate the configured MCP host allow-list into SDK transport settings. Without this the SDK infers its own allow-list from the bind host, which defaults to 127.0.0.1 and therefore rejects every remote client with HTTP 421 — fatal for a remotely deployed backend. """ allowed = _parse_allowed_hosts() if "*" in allowed: # Explicit, logged opt-out. Kept as an escape hatch for environments # behind a proxy that rewrites Host unpredictably, but never the default. logger.warning( "MCP DNS-rebinding protection is disabled (mcp_allowed_hosts='*'). " "Set MCP_ALLOWED_HOSTS to the real deployment host(s) instead." ) return TransportSecuritySettings(enable_dns_rebinding_protection=False) return TransportSecuritySettings( enable_dns_rebinding_protection=True, allowed_hosts=allowed, # Browser clients send Origin; reuse the already-maintained CORS list so # there is one place to declare trusted web origins. Non-browser MCP # clients send no Origin at all, which the SDK treats as allowed. allowed_origins=[o.strip() for o in settings.cors_allow_origins.split(",") if o.strip()], ) def build_mcp_asgi_app() -> ASGIApp: """Return the Streamable HTTP ASGI app for the MCP server, auth-guarded. streamable_http_path="/" is required here: MCPServer.streamable_http_app() registers its own internal route at "/mcp" by default, and this app is itself mounted at "/mcp" in api/main.py — without overriding the internal path to "/", the effective external path would be the confusing "/mcp/mcp" instead of "/mcp". """ asgi_app = mcp.streamable_http_app( streamable_http_path="/", transport_security=_build_transport_security(), ) asgi_app.add_middleware(MCPAuthMiddleware) return asgi_app async def get_mcp_status(public_url: str) -> dict: """Assemble the MCP status payload shown on the System Status page. Owned by this module rather than the status route so that MCP internals (the tool registry, the allow-list format, the stats tracker) stay behind one boundary; the route only supplies public_url, which is the one value only the HTTP layer can know. """ stats = get_mcp_stats_tracker().snapshot() # list_tools() reads the in-memory registry populated by @mcp.tool() at # import time, so the panel always reflects what is actually advertised # rather than a hand-maintained duplicate list. tools = await mcp.list_tools() return { "endpoint_url": public_url, "auth_required": settings.auth_enabled, "allowed_hosts": _parse_allowed_hosts(), "tools": [ { "name": tool.name, "description": (tool.description or "").strip().split("\n")[0], "calls": entry.calls if entry else 0, "errors": entry.errors if entry else 0, "avg_duration_ms": entry.avg_duration_ms if entry else None, "last_called_at": ( entry.last_called_at.isoformat() if entry and entry.last_called_at else None ), } for tool, entry in ((tool, stats.get(tool.name)) for tool in tools) ], }